Add the internal _feature_version to compile()

gaogaotiantian · Jul 31, 2024 · 5bf6c22 · 5bf6c22
1 parent c31bd9b
commit 5bf6c22
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/src/viztracer/code_monkey.py b/src/viztracer/code_monkey.py
@@ -347,14 +347,15 @@ def add_instrument(self, inst_type: str, inst_args: Dict[str, Dict]) -> None:
     def add_source_processor(self):
         self.source_processor = SourceProcessor()
 
-    def compile(self, source, filename, mode, flags=0, dont_inherit=False, optimize=-1):
+    def compile(self, source, filename, mode, flags=0, dont_inherit=False, optimize=-1, *, _feature_version=-1):
         if self.source_processor is not None:
             source = self.source_processor.process(source)
         if self.ast_transformers:
-            tree = self._compile(source, filename, mode, flags | ast.PyCF_ONLY_AST, dont_inherit, optimize)
+            tree = self._compile(source, filename, mode, flags | ast.PyCF_ONLY_AST,
+                                 dont_inherit, optimize, _feature_version=_feature_version)
             for trans in self.ast_transformers:
                 trans.visit(tree)
                 ast.fix_missing_locations(tree)
-            return self._compile(tree, filename, mode, flags, dont_inherit, optimize)
+            return self._compile(tree, filename, mode, flags, dont_inherit, optimize, _feature_version=_feature_version)
 
-        return self._compile(source, filename, mode, flags, dont_inherit, optimize)
+        return self._compile(source, filename, mode, flags, dont_inherit, optimize, _feature_version=_feature_version)
diff --git a/tests/test_cmdline.py b/tests/test_cmdline.py
@@ -7,6 +7,7 @@
 import re
 import signal
 import sys
+import textwrap
 import tempfile
 from contextlib import contextmanager
 from unittest.case import skipIf
@@ -400,6 +401,15 @@ def test_log_var(self):
                       expected_output_file="result.json",
                       expected_entries=12)
 
+        code_ast = textwrap.dedent("""
+            import ast
+            tree = ast.parse("a = 1")
+        """)
+        self.template(["viztracer", "--log_var", "tree", "-o", "result.json", "cmdline_test.py"],
+                      script=code_ast,
+                      expected_output_file="result.json",
+                      expected_entries=4)
+
     def test_log_attr(self):
         self.template(["viztracer", "--log_attr", "a.*", "-o", "result.json", "cmdline_test.py"],
                       script=file_log_attr,