diff --git a/src/smolagents/local_python_executor.py b/src/smolagents/local_python_executor.py index 9477be9..e46d87a 100644 --- a/src/smolagents/local_python_executor.py +++ b/src/smolagents/local_python_executor.py @@ -641,11 +641,13 @@ def evaluate_subscript( return value[index] elif index in value: return value[index] - elif isinstance(index, str) and isinstance(value, Mapping): - close_matches = difflib.get_close_matches(index, list(value.keys())) - if len(close_matches) > 0: - return value[close_matches[0]] - raise InterpreterError(f"Could not index {value} with '{index}'.") + else: + error_message = f"Could not index {value} with '{index}'." + if isinstance(index, str) and isinstance(value, Mapping): + close_matches = difflib.get_close_matches(index, list(value.keys())) + if len(close_matches) > 0: + error_message += f" Maybe you meant one of these indexes instead: {str(close_matches)}" + raise InterpreterError(error_message) def evaluate_name( diff --git a/tests/test_python_interpreter.py b/tests/test_python_interpreter.py index 0e83177..540720f 100644 --- a/tests/test_python_interpreter.py +++ b/tests/test_python_interpreter.py @@ -897,3 +897,11 @@ shift_intervals code = "import doctest;doctest.inspect.os.system('echo bad command passed')" with pytest.raises(AttributeError): evaluate_python_code(code, authorized_imports=["doctest"]) + + def test_close_matches_subscript(self): + code = 'capitals = {"Czech Republic": "Prague", "Monaco": "Monaco", "Bhutan": "Thimphu"};capitals["Butan"]' + with pytest.raises(Exception) as e: + evaluate_python_code(code) + assert "Maybe you meant one of these indexes instead" in str( + e + ) and "['Bhutan']" in str(e).replace("\\", "")