from evaluator import *
DESCRIPTION = "Test if the model can shorten a line of python with an equal line."
TAGS = ['code', 'python']
question = """
Shorten this python code. Keep the name the same.
```
def foo(arr, keys):
return [[x[y] if y in x else -5 for y in keys] for x in arr]
```
"""
test_case, answer = make_python_test([("foo([{1: 2, 3: 4, 5: 6}], [1, 2, 3, 4])", "[[2, -5, 4, -5]]")])
TestShortenPyGet = question >> LLMRun() >> (SubstringEvaluator(".get") & (ExtractCode() >> PythonRun(test_case) >> SubstringEvaluator(answer)))
if __name__ == "__main__":
print(run_test(TestShortenPyGet))