From 5bf6c22f6b134302e5f83ba38aa281837d8fa267 Mon Sep 17 00:00:00 2001
From: Tian Gao <gaogaotiantian@hotmail.com>
Date: Tue, 30 Jul 2024 18:40:10 -0700
Subject: [PATCH] Add the internal _feature_version to compile()

---
 src/viztracer/code_monkey.py |  9 +++++----
 tests/test_cmdline.py        | 10 ++++++++++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/viztracer/code_monkey.py b/src/viztracer/code_monkey.py
index c8b46ad..68e382e 100644
--- 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
index 777ccb2..6f994e6 100644
--- 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,