Return()

tests/test_definitions.py

@@ -11,3 +11,26 @@ def test_function_def():
     def target():
         pass
 ''')
+
+
+def test_return():
+
+    with closing(StringIO()) as f:
+
+        @Commentor(f)
+        def target():
+            a = 1
+            return a + 1
+
+        assert target() == 2
+
+        asserteq_or_print(
+            f.getvalue(), '''
+        def target():
+            a = 1
+            return a + 1
+            """
+            1 : a
+            2 : a + 1
+            """
+''')

tests/test_utils.py

@@ -1,7 +1,11 @@
+from io import StringIO
+from contextlib import closing
 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"])
+        value = value.strip("\n").rstrip(" ")
+        ground_truth = ground_truth.strip("\n").rstrip(" ")
+        assert value == ground_truth, "\n".join(["\n\n<<<<<<<< VALUE", value, "========================", ground_truth, ">>>>>>>> GROUND\n"])