From 499a8ad1d3f3d3954492be9877abc7906887cce4 Mon Sep 17 00:00:00 2001 From: gaogaotiantian Date: Sun, 28 Nov 2021 19:31:27 -0800 Subject: [PATCH] Add conditional magic comment (#185) * Add conditional magic comment --- docs/source/custom_event_intro.rst | 13 ++++++++++++- src/viztracer/code_monkey.py | 20 ++++++++++++++++++-- src/viztracer/viztracer.py | 18 ++++++++++-------- tests/test_codemonkey.py | 21 ++++++++++++++------- 4 files changed, 54 insertions(+), 18 deletions(-) diff --git a/docs/source/custom_event_intro.rst b/docs/source/custom_event_intro.rst index 6ac0f095..ed03d857 100644 --- a/docs/source/custom_event_intro.rst +++ b/docs/source/custom_event_intro.rst @@ -53,7 +53,7 @@ viztracer to trace it # !viztracer: log_instant("start logging") a = 3 - # !viztracer: log_var("a", 3) + # !viztracer: log_var("a", a) Or you can use inline magic comment ``# !viztracer: log``, which will log the assigned value if the statement is an assign or it will log an instant event indicating this line is executed @@ -66,6 +66,17 @@ or it will log an instant event indicating this line is executed # This will log the variable a a = 3 # !viztracer: log +You can also do conditional log with ``if`` + +.. code-block:: python + + # This will log the variable a + a = 3 # !viztracer: log if a == 3 + # This has the same effect + # !viztracer: log_var("a", a) + +You need ``--magic_comment`` option for ``viztracer`` to trigger the magic comment + .. code-block:: viztracer --magic_comment your_program.py diff --git a/src/viztracer/code_monkey.py b/src/viztracer/code_monkey.py index 82957fbb..fa7863be 100644 --- a/src/viztracer/code_monkey.py +++ b/src/viztracer/code_monkey.py @@ -292,8 +292,14 @@ class SourceProcessor: """ def __init__(self): self.re_patterns = [ - (re.compile(r"(\s*)#\s*!viztracer:\s*(.*?)$"), self.line_transform), - (re.compile(r"(.*\S.*)#\s*!viztracer:\s*log$"), self.inline_transform) + # !viztracer: log_var("var", var) + (re.compile(r"(\s*)#\s*!viztracer:\s*(log_.*?\(.*\))\s*$"), self.line_transform), + # a = 3 # !viztracer: log + (re.compile(r"(.*\S.*)#\s*!viztracer:\s*log\s*$"), self.inline_transform), + # !viztracer: log_var("var", var) if var > 3 + (re.compile(r"(\s*)#\s*!viztracer:\s*(log_.*?\(.*\))\s*if\s+(.*?)\s*$"), self.line_transform_condition), + # a = 3 # !viztracer: log if a != 3 + (re.compile(r"(.*\S.*)#\s*!viztracer:\s*log\s*if\s+(.*?)\s*$"), self.inline_transform_condition) ] def process(self, source: Any): @@ -318,6 +324,9 @@ def process(self, source: Any): def line_transform(self, re_match): return f"{re_match.group(1)}__viz_tracer__.{re_match.group(2)}" + def line_transform_condition(self, re_match): + return f"{re_match.group(1)}if {re_match.group(3)}: __viz_tracer__.{re_match.group(2)};" + def inline_transform(self, re_match): stmt = re_match.group(1) if "=" in stmt: @@ -325,6 +334,13 @@ def inline_transform(self, re_match): return f"{stmt}; __viz_tracer__.log_var('{val_assigned}', ({val_assigned}))" return f"{stmt}; __viz_tracer__.log_instant('{stmt.strip()}')" + def inline_transform_condition(self, re_match): + stmt = re_match.group(1) + if "=" in stmt: + val_assigned = stmt[:stmt.index("=")].strip() + return f"{stmt}; __viz_tracer__.log_var('{val_assigned}', ({val_assigned}), cond={re_match.group(2)})" + return f"{stmt}; __viz_tracer__.log_instant('{stmt.strip()}', cond={re_match.group(2)});" + class CodeMonkey: def __init__(self, file_name: str) -> None: diff --git a/src/viztracer/viztracer.py b/src/viztracer/viztracer.py index e260ca83..73a0268e 100644 --- a/src/viztracer/viztracer.py +++ b/src/viztracer/viztracer.py @@ -141,14 +141,16 @@ def signal_stop(signum, frame): signal.signal(signal.SIGUSR1, signal_start) signal.signal(signal.SIGUSR2, signal_stop) - def log_instant(self, name: str, args: Any = None, scope: str = "t") -> None: - self.add_instant(name, args=args, scope=scope) - - def log_var(self, name, var) -> None: - if isinstance(var, (float, int)): - self.add_counter(name, {"value": var}) - else: - self.add_instant(name, args={"object": objprint.objstr(var, color=False)}, scope="t") + def log_instant(self, name: str, args: Any = None, scope: str = "t", cond: bool = True) -> None: + if cond: + self.add_instant(name, args=args, scope=scope) + + def log_var(self, name, var, cond: bool = True) -> None: + if cond: + if isinstance(var, (float, int)): + self.add_counter(name, {"value": var}) + else: + self.add_instant(name, args={"object": objprint.objstr(var, color=False)}, scope="t") def log_event(self, event_name: str) -> VizEvent: call_frame = sys._getframe(1) diff --git a/tests/test_codemonkey.py b/tests/test_codemonkey.py index cc09db32..efe23b9a 100644 --- a/tests/test_codemonkey.py +++ b/tests/test_codemonkey.py @@ -80,26 +80,33 @@ def test_get_assign_log_nodes(self): file_magic_comment = """ +def test(): + pass # !viztracer: log_instant("test") a = 3 # !viztracer: log # !viztracer: log_var("a", a) +# !viztracer: log_var("a", a) if a == 3 +# !viztracer: log_var("a", a) if a != 3 +a = 3 # !viztracer: log if a == 3 +a = 3 # !viztracer: log if a != 3 +test() # !viztracer: log if a == 3 """ class TestMagicComment(CmdlineTmpl): def test_log_var(self): def check_func(data): - found_instant = False - found_var = False + instant_count = 0 + var_count = 0 for event in data["traceEvents"]: if event["ph"] == "i": - self.assertEqual(event["name"], "test") - found_instant = True + self.assertIn("test", event["name"]) + instant_count += 1 elif event["ph"] == "C": self.assertEqual(event["name"], "a") self.assertEqual(event["args"], {"value": 3}) - found_var = True - self.assertTrue(found_instant) - self.assertTrue(found_var) + var_count += 1 + self.assertEqual(instant_count, 2) + self.assertEqual(var_count, 4) self.template(["viztracer", "--magic_comment", "cmdline_test.py"], script=file_magic_comment, check_func=check_func)