From 5017ce273634d876970d712fb2a927ec13e829f3 Mon Sep 17 00:00:00 2001
From: Yuyao Huang <huangyuyao@outlook.com>
Date: Fri, 19 Apr 2024 14:54:14 +0800
Subject: [PATCH] automatic regressive testing

---
 tests/test_definitions.py               |  9 +++++---
 tests/test_expressions.py               | 29 +++++++++++++++++++++++--
 tests/test_literals.py                  |  8 +++++--
 tests/test_utils.py                     |  7 ++++++
 trace_commentor/flags.py                |  7 +++++-
 trace_commentor/handlers/definitions.py |  2 +-
 trace_commentor/handlers/statements.py  |  2 +-
 trace_commentor/utils.py                |  5 +++--
 8 files changed, 57 insertions(+), 12 deletions(-)
 create mode 100644 tests/test_utils.py

diff --git a/tests/test_definitions.py b/tests/test_definitions.py
index 39c1127..c7777af 100644
--- a/tests/test_definitions.py
+++ b/tests/test_definitions.py
@@ -1,4 +1,4 @@
-from trace_commentor import Commentor
+from test_utils import *
 
 
 def test_function_def():
@@ -6,5 +6,8 @@ def test_function_def():
     @Commentor()
     def target():
         pass
-    
-    print(target())
+
+    asserteq_or_print(target(), '''
+    def target():
+        pass
+''')
diff --git a/tests/test_expressions.py b/tests/test_expressions.py
index b787721..38d672c 100644
--- a/tests/test_expressions.py
+++ b/tests/test_expressions.py
@@ -1,4 +1,5 @@
 from trace_commentor import Commentor
+from test_utils import asserteq_or_print
 
 
 def test_binop():
@@ -6,5 +7,29 @@ def test_binop():
     @Commentor()
     def target():
         1 + 1
-    
-    print(target())
+
+    asserteq_or_print(
+        target(), '''
+    def target():
+        1 + 1
+        """
+        2 : 1 + 1
+        """
+''')
+
+
+def test_binop_cascade():
+
+    @Commentor()
+    def target():
+        1 + 1 + 1
+
+    asserteq_or_print(
+        target(), '''
+    def target():
+        1 + 1 + 1
+        """
+        2 : 1 + 1
+        3 : 1 + 1 + 1
+        """
+''')
diff --git a/tests/test_literals.py b/tests/test_literals.py
index d3ef152..1c95cfb 100644
--- a/tests/test_literals.py
+++ b/tests/test_literals.py
@@ -1,4 +1,4 @@
-from trace_commentor import Commentor
+from test_utils import *
 
 
 def test_constant():
@@ -7,4 +7,8 @@ def test_constant():
     def target():
         1
     
-    print(target())
+    asserteq_or_print(target(),
+    '''
+    def target():
+        1
+''')
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..b6b9dad
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,7 @@
+from trace_commentor import flags, Commentor
+
+def asserteq_or_print(value, ground_truth):
+    if flags.DEBUG or flags.PRINT:
+        print(value)
+    else:
+        assert value == ground_truth.strip("\n"), "\n".join(["\n\n<<<<<<<< VALUE", value, "========================", ground_truth.strip("\n"), ">>>>>>>> GROUND\n"])
diff --git a/trace_commentor/flags.py b/trace_commentor/flags.py
index e2c4948..61b02b5 100644
--- a/trace_commentor/flags.py
+++ b/trace_commentor/flags.py
@@ -1,4 +1,9 @@
-DEBUG = False
+import os
+
+bool_env = lambda name: os.environ.get(name, "false").lower() in ('true', '1', 'yes')
+
+DEBUG = bool_env("DEBUG")
+PRINT = bool_env("PRINT")
 INDENT = 4
 SOURCE = 1
 COMMENT = 2
diff --git a/trace_commentor/handlers/definitions.py b/trace_commentor/handlers/definitions.py
index 0632c2e..fcf7dfd 100644
--- a/trace_commentor/handlers/definitions.py
+++ b/trace_commentor/handlers/definitions.py
@@ -2,7 +2,7 @@ from .. import flags
 from ..utils import to_source
 
 def FunctionDef(self, cmtor):
-    cmtor.append_source("def function():")
+    cmtor.append_source(f"def {self.name}():")
     cmtor.indent += flags.INDENT
     
     for stmt in self.body:
diff --git a/trace_commentor/handlers/statements.py b/trace_commentor/handlers/statements.py
index a7e1ee2..cf7c449 100644
--- a/trace_commentor/handlers/statements.py
+++ b/trace_commentor/handlers/statements.py
@@ -1,2 +1,2 @@
 def Pass(self, cmtor):
-    cmtor.append("pass")
+    pass
diff --git a/trace_commentor/utils.py b/trace_commentor/utils.py
index 10891bc..37da372 100644
--- a/trace_commentor/utils.py
+++ b/trace_commentor/utils.py
@@ -1,3 +1,4 @@
+import os
 import astor
 import inspect
 from . import flags
@@ -5,8 +6,8 @@ from . import flags
 def sign(line: str, depth=1):
     if flags.DEBUG:
         currentframe = inspect.currentframe()
-        outerframe = inspect.getouterframes(currentframe, depth)
-        debug_msg = outerframe[1][3]
+        outerframe = inspect.getouterframes(currentframe, 1)[depth]
+        debug_msg = f"{outerframe.function} @ {os.path.relpath(outerframe.filename)}:{outerframe.lineno}"
         return f"{line}  ---  {debug_msg}"
     else:
         return line