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,