Assign();Name();Call()

tests/test_definitions.py

@@ -3,7 +3,7 @@ from test_utils import *
 
 def test_function_def():
 
-    @Commentor()
+    @Commentor("<return>")
     def target():
         pass
 

tests/test_expressions.py

@@ -4,7 +4,7 @@ from test_utils import asserteq_or_print
 
 def test_binop():
 
-    @Commentor()
+    @Commentor("<return>")
     def target():
         1 + 1
 
@@ -20,7 +20,7 @@ def test_binop():
 
 def test_binop_cascade():
 
-    @Commentor()
+    @Commentor("<return>")
     def target():
         1 + 1 + 1
 
@@ -33,3 +33,19 @@ def test_binop_cascade():
         3 : 1 + 1 + 1
         """
 ''')
+
+
+def test_call_print():
+
+    @Commentor("<return>")
+    def target():
+        print("This line will be printed.")
+
+    asserteq_or_print(
+        target(), '''
+    def target():
+        print('This line will be printed.')
+        """
+        None : print('This line will be printed.')
+        """
+''')

tests/test_literals.py

@@ -3,12 +3,11 @@ from test_utils import *
 
 def test_constant():
 
-    @Commentor()
+    @Commentor("<return>")
     def target():
         1
-    
-    asserteq_or_print(target(),
-    '''
+
+    asserteq_or_print(target(), '''
     def target():
         1
 ''')

tests/test_torch.py

@@ -1,37 +0,0 @@
-import torch
-import torch.nn as nn
-
-from trace_commentor.parser import *
-
-
-def test_func_no_arg():
-
-    @analyse
-    def target():
-
-        x = torch.ones(4, 5)
-        for i in range(3):
-            x = x[..., None, :]
-
-        a = torch.randn(309, 110, 3)[:100]
-        f = nn.Linear(3, 128)
-        b = f(a.reshape(-1, 3)).reshape(309, 110, 128)
-        c = torch.concat((a, b), dim=-1)
-
-        return c.flatten()
-    
-    print()
-    target()
-
-
-def test_for_loop():
-
-    @analyse
-    def target():
-        a = 1
-        for i in range(3):
-            a += 1
-        print(a)
-    
-    print()
-    target()

tests/test_variables.py

@@ -0,0 +1,20 @@
+from test_utils import *
+
+
+def test_assign():
+
+    @Commentor("<return>")
+    def target():
+        myint = 7
+        print(myint)
+
+    asserteq_or_print(
+        target(), '''
+    def target():
+        myint = 7
+        print(myint)
+        """
+        7 : myint
+        None : print(myint)
+        """
+''')