Add conditional magic comment (#185)

* Add conditional magic comment

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

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,13 +324,23 @@ 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:
             val_assigned = stmt[:stmt.index("=")].strip()
             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:

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)

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)