Add test for from __future__ import

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

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 == []