Python: add support for auxiliary code (like for Prolog)

1 files changed, 4 insertions, 4 deletions

diff --git a/python/problems/while_and_if/buy_five/common.py b/python/problems/while_and_if/buy_five/common.py
index 0ec7238..ef85b07 100644
--- a/python/problems/while_and_if/buy_five/common.py
+++ b/python/problems/while_and_if/buy_five/common.py
@@ -23,7 +23,7 @@ hint_type = {
     'nonumber': Hint('nonumber'),
 }
 
-def test(python, code):
+def test(python, code, aux_code=''):
     # List of inputs: (expression to eval, stdin).
     test_in = [
         (None, '1\n1\n1\n1\n1\n'),
@@ -42,7 +42,7 @@ def test(python, code):
     ]
 
     # List of outputs: (expression result, stdout, stderr, exception).
-    answers = python(code=code, inputs=test_in, timeout=1.0)
+    answers = python(code=aux_code+code, inputs=test_in, timeout=1.0)
     outputs = [ans[1] for ans in answers]
 
     n_correct = 0
@@ -60,10 +60,10 @@ def test(python, code):
         hints.append({'id': 'problematic_test_case', 'args': {'testin': str(tin), 'testout': str(tout)}})
     return passed, hints
 
-def hint(python, code):
+def hint(python, code, aux_code=''):
     # run one test first to see if there are any exceptions
     test_in = [(None, '1\n1\n1\n1\n1\n')]
-    answer = python(code=code, inputs=test_in, timeout=1.0)
+    answer = python(code=aux_code+code, inputs=test_in, timeout=1.0)
     exc = get_exception_desc(answer[0][3])
     if exc: return exc