diff --git a/openc3-cosmos-script-runner-api/scripts/script_instrumentor.py b/openc3-cosmos-script-runner-api/scripts/script_instrumentor.py index 559456509..5ce37ac87 100644 --- a/openc3-cosmos-script-runner-api/scripts/script_instrumentor.py +++ b/openc3-cosmos-script-runner-api/scripts/script_instrumentor.py @@ -155,6 +155,12 @@ def track_reached(self, node): ast.copy_location(if_node, node) return if_node + def track_import_from(self, node): + # Don't tract from __future__ imports because they must come first or: + # SyntaxError: from __future__ imports must occur at the beginning of the file + if node.module != '__future__': + return self.track_enter_leave(node) + # Notes organized (including newlines) per https://docs.python.org/3/library/ast.html#abstract-grammar # Nodes that change control flow are processed by track_reached, otherwise we track_enter_leave visit_FunctionDef = track_reached @@ -187,7 +193,7 @@ def track_reached(self, node): visit_Assert = track_enter_leave visit_Import = track_enter_leave - visit_ImportFrom = track_enter_leave + visit_ImportFrom = track_import_from visit_Global = track_enter_leave visit_Nonlocal = track_enter_leave diff --git a/openc3-cosmos-script-runner-api/test/scripts/test_script_instrumentor.py b/openc3-cosmos-script-runner-api/test/scripts/test_script_instrumentor.py index 8ed76c87f..13fe8e7b4 100644 --- a/openc3-cosmos-script-runner-api/test/scripts/test_script_instrumentor.py +++ b/openc3-cosmos-script-runner-api/test/scripts/test_script_instrumentor.py @@ -147,7 +147,6 @@ def test_match_script(mock_running_script): """ parsed = ast.parse(script) tree = ScriptInstrumentor("testfile.py").visit(parsed) - print(ast.dump(tree, indent=4)) compiled = compile(tree, filename="testfile.py", mode="exec") exec(compiled, {"RunningScript": mock_running_script}) @@ -174,7 +173,6 @@ def test_exception_script(mock_running_script): """ parsed = ast.parse(script) tree = ScriptInstrumentor("testfile.py").visit(parsed) - print(ast.dump(tree, indent=4)) compiled = compile(tree, filename="testfile.py", mode="exec") exec(compiled, {"RunningScript": mock_running_script}) @@ -191,3 +189,19 @@ def test_exception_script(mock_running_script): assert mock_running_script.exceptions == [ ("testfile.py", 3), ] + + +def test_import_future_script(mock_running_script): + script = "from __future__ import annotations\nprint('hi')" + parsed = ast.parse(script) + tree = ScriptInstrumentor("testfile.py").visit(parsed) + compiled = compile(tree, filename="testfile.py", mode="exec") + exec(compiled, {"RunningScript": mock_running_script}) + + assert mock_running_script.pre_lines == [ + ("testfile.py", 2), + ] + assert mock_running_script.post_lines == [ + ("testfile.py", 2), + ] + assert mock_running_script.exceptions == []