For(); List()

.gitignore

@@ -1,3 +1,4 @@
 __pycache__
 *.log
+*.log.py
 .vscode-upload.json

experiment.py

@@ -5,20 +5,12 @@ def test():
 
     @Commentor()
     def target():
-        x = 2
+        # return only odd numbers - 1,3,5,7,9
-        if x > 3:
+        for x in range(10):
-            x = 2 * x
+            # Check if x is even
-            y = 1
+            if x % 2 == 0:
-        elif x > 2:
+                continue
-            x = 4 * x
+            print(x)
-            y = 2
-        elif x > 3:
-            x = 4 * x
-            y = 3
-        else:
-            x = 8 * x
-            y = 5
-
 
     print(target())
 

tests/test_conditions.py

@@ -1,61 +0,0 @@
-from test_utils import *
-
-
-def test_constant():
-
-    @Commentor("<return>")
-    def target():
-        x = 2
-        print(x == 2)
-
-    asserteq_or_print(target(), '''
-    def target():
-        x = 2
-        print(x == 2)
-        """
-        True : x == 2
-        None : print(x == 2)
-        """
-''')
-
-
-def test():
-
-    @Commentor("<return>")
-    def target():
-        x = 2
-        if x > 3:
-            x = 2 * x
-            y = 1
-        elif x > 2:
-            x = 4 * x
-            y = 2
-        elif x > 3:
-            x = 4 * x
-            y = 3
-        else:
-            x = 8 * x
-            y = 5
-    
-    asserteq_or_print(target(), '''
-    def target():
-        x = 2
-        if x > 3:  # False
-            x = 2 * x
-            y = 1
-        elif x > 2:  # False
-            x = 4 * x
-            y = 2
-        elif x > 3:  # False
-            x = 4 * x
-            y = 3
-        else:  # True
-            x = 8 * x
-            """
-            2 : x
-            16 : 8 * x
-            ----------
-            16 : x
-            """
-            y = 5
-''')

tests/test_control_flow.py

@@ -0,0 +1,208 @@
+from test_utils import *
+
+
+def test_constant():
+
+    @Commentor("<return>")
+    def target():
+        x = 2
+        print(x == 2)
+
+    asserteq_or_print(target(), '''
+    def target():
+        x = 2
+        print(x == 2)
+        """
+        True : x == 2
+        None : print(x == 2)
+        """
+''')
+
+
+def test_if():
+
+    @Commentor("<return>")
+    def target():
+        x = 2
+        if x > 3:
+            x = 2 * x
+            y = 1
+        elif x > 2:
+            x = 4 * x
+            y = 2
+        elif x > 3:
+            x = 4 * x
+            y = 3
+        else:
+            x = 8 * x
+            y = 5
+    
+    asserteq_or_print(target(), '''
+    def target():
+        x = 2
+        if x > 3:  # False
+            x = 2 * x  # skipped
+            y = 1  # skipped
+        elif x > 2:  # False
+            x = 4 * x  # skipped
+            y = 2  # skipped
+        elif x > 3:  # False
+            x = 4 * x  # skipped
+            y = 3  # skipped
+        else:    # True
+            x = 8 * x
+            """
+            2 : x
+            16 : 8 * x
+            ----------
+            16 : x
+            """
+            y = 5
+''')
+
+
+def test_for():
+    
+    with closing(StringIO()) as f:
+        @Commentor(f)
+        def target():
+            odds = []
+            # return only odd numbers - 1,3,5,7,9
+            for x in range(10):
+                # Check if x is even
+                if x % 2 == 0:
+                    continue
+                odds.append(x)
+            return odds
+        
+        assert target() == [1, 3, 5, 7, 9]
+        asserteq_or_print(f.getvalue(), '''
+        def target():
+            odds = []
+            for x in range(10):
+                
+                ###### !new iteration! ######
+                """
+                0 : __REG__for_loop_iter_once
+                ----------
+                0 : x
+                """
+                # if x % 2 == 0:  # True
+                #     continue  # True
+                # odds.append(x)  # skipped
+                
+                ###### !new iteration! ######
+                """
+                1 : __REG__for_loop_iter_once
+                ----------
+                1 : x
+                """
+                # if x % 2 == 0:  # False
+                #     continue  # skipped
+                # odds.append(x)
+                """
+                1 : x
+                None : odds.append(x)
+                """
+                
+                ###### !new iteration! ######
+                """
+                2 : __REG__for_loop_iter_once
+                ----------
+                2 : x
+                """
+                # if x % 2 == 0:  # True
+                #     continue  # True
+                # odds.append(x)  # skipped
+                
+                ###### !new iteration! ######
+                """
+                3 : __REG__for_loop_iter_once
+                ----------
+                3 : x
+                """
+                # if x % 2 == 0:  # False
+                #     continue  # skipped
+                # odds.append(x)
+                """
+                3 : x
+                None : odds.append(x)
+                """
+                
+                ###### !new iteration! ######
+                """
+                4 : __REG__for_loop_iter_once
+                ----------
+                4 : x
+                """
+                # if x % 2 == 0:  # True
+                #     continue  # True
+                # odds.append(x)  # skipped
+                
+                ###### !new iteration! ######
+                """
+                5 : __REG__for_loop_iter_once
+                ----------
+                5 : x
+                """
+                # if x % 2 == 0:  # False
+                #     continue  # skipped
+                # odds.append(x)
+                """
+                5 : x
+                None : odds.append(x)
+                """
+                
+                ###### !new iteration! ######
+                """
+                6 : __REG__for_loop_iter_once
+                ----------
+                6 : x
+                """
+                # if x % 2 == 0:  # True
+                #     continue  # True
+                # odds.append(x)  # skipped
+                
+                ###### !new iteration! ######
+                """
+                7 : __REG__for_loop_iter_once
+                ----------
+                7 : x
+                """
+                # if x % 2 == 0:  # False
+                #     continue  # skipped
+                # odds.append(x)
+                """
+                7 : x
+                None : odds.append(x)
+                """
+                
+                ###### !new iteration! ######
+                """
+                8 : __REG__for_loop_iter_once
+                ----------
+                8 : x
+                """
+                # if x % 2 == 0:  # True
+                #     continue  # True
+                # odds.append(x)  # skipped
+                
+                ###### !new iteration! ######
+                """
+                9 : __REG__for_loop_iter_once
+                ----------
+                9 : x
+                """
+                if x % 2 == 0:  # False
+                    continue  # skipped
+                odds.append(x)
+                """
+                9 : x
+                None : odds.append(x)
+                """
+            return odds
+            """
+            [1, 3, 5, 7, 9] : odds
+            """''')
+        
+        

trace_commentor/commentor.py

@@ -11,6 +11,7 @@ from . import handlers
 from . import formatters
 from . import flags
 from .utils import sign, to_source
+from rich.syntax import Syntax
 
 
 class Commentor(object):
@@ -21,9 +22,11 @@ class Commentor(object):
         self._return = None
         self._formatters = fmt + formatters.LIST
         self._lines = []
+        self._lines_category = []
         self.indent = 0
         self.state = flags.SOURCE
         self.file = output
+        self._stack_event = flags.NORMAL
 
     def __call__(self, func):
 
@@ -49,8 +52,16 @@ class Commentor(object):
             if self.file == "<return>":
                 return code
             elif self.file == "<stderr>":
+                if sys.stderr.isatty():
+                    syntax = Syntax(code, "python")
+                    rich.print(syntax, file=sys.stderr)
+                else:
                     rich.print(code, file=sys.stderr)
             elif self.file == "<stdout>":
+                if sys.stdout.isatty():
+                    syntax = Syntax(code, "python")
+                    rich.print(syntax, file=sys.stdout)
+                else:
                     rich.print(code, file=sys.stdout)
             elif isinstance(self.file, IOBase):
                 rich.print(code, file=self.file)
@@ -93,19 +104,28 @@ class Commentor(object):
         else:
             return repr
 
+    def next_line(self) -> int:
+        return len(self._lines)
+
     def __append(self, line):
         self._lines.append(" " * self.indent + str(line))
+        self._lines_category.append((self.state, self.indent))
+        return len(self._lines) - 1
 
     def append_source(self, line=None):
         if self.state == flags.COMMENT:
             self.__append('"""')
             self.state = flags.SOURCE
         if line is not None:
-            self.__append(sign(line, 2))
+            return self.__append(sign(line, 2))
 
     def append_comment(self, line=None):
         if self.state == flags.SOURCE:
-            self.__append('"""')
             self.state = flags.COMMENT
+            self.__append('"""')
         if line is not None:
-            self.__append(sign(line, 2))
+            return self.__append(sign(line, 2))
+
+    def typeset(self):
+        return "\n".join(
+            (c[1] * " " + l for l, c in zip(self._lines, self._lines_category)))

trace_commentor/flags.py

@@ -1,3 +1,4 @@
+import ast
 import os
 
 bool_env = lambda name: os.environ.get(name, "false").lower() in ('true', '1', 'yes')
@@ -7,3 +8,16 @@ PRINT = bool_env("PRINT")
 INDENT = 4
 SOURCE = 1
 COMMENT = 2
+NORMAL = 0
+BREAK = 8
+CONTINUE = 16
+
+REG = lambda i: f"__REG{i}"
+
+APPEND_SOURCE_BY_THEMSELVES = [
+    ast.If, ast.For
+]
+
+ASSIGN_SILENT = [
+    ast.Constant, ast.List
+]

trace_commentor/handlers/__init__.py

@@ -1,6 +1,6 @@
 from .definitions import FunctionDef, Return
 from .statements import Pass, Assign
 from .expressions import Expr, BinOp, Call, Compare
-from .literals import Constant, Tuple
+from .literals import Constant, Tuple, List
 from .variables import Name
-from .control_flow import If
+from .control_flow import If, For, Continue, Break

trace_commentor/handlers/control_flow.py

@@ -1,7 +1,6 @@
 import ast
 from .. import flags
-from ..utils import to_source, APPEND_SOURCE_BY_THEMSELVES
+from ..utils import to_source
-
 
 ELIF = 2
 PASS = 4
@@ -25,10 +24,13 @@ def If(self, cmtor, state=0):
 
     cmtor.indent += flags.INDENT
     for stmt in self.body:
-        if type(stmt) not in APPEND_SOURCE_BY_THEMSELVES:
+        if type(stmt) not in flags.APPEND_SOURCE_BY_THEMSELVES:
             cmtor.append_source(to_source(stmt))
         if test:
             cmtor.process(stmt)
+            test = cmtor._stack_event == flags.NORMAL
+        else:
+            cmtor._lines[-1] += "  # skipped"
     cmtor.append_source()
     cmtor.indent -= flags.INDENT
 
@@ -42,9 +44,67 @@ def If(self, cmtor, state=0):
 
             cmtor.indent += flags.INDENT
             for stmt in self.orelse:
-                if type(stmt) not in APPEND_SOURCE_BY_THEMSELVES:
+                if type(stmt) not in flags.APPEND_SOURCE_BY_THEMSELVES:
                     cmtor.append_source(to_source(stmt))
                 if test:
                     cmtor.process(stmt)
+                    test = cmtor._stack_event == flags.NORMAL
             cmtor.append_source()
             cmtor.indent -= flags.INDENT
+
+
+def For(self, cmtor):
+    cmtor.append_source(to_source(ast.For(self.target, self.iter, [], [])))
+
+    loop_start: int = cmtor.next_line()
+
+    cmtor.indent += flags.INDENT
+    self_indent = cmtor.indent
+
+    REG_it = flags.REG("__for_loop_iter_once")
+    iter_obj = cmtor.eval(self.iter, format=False)
+    for it in iter_obj:
+        last_iter_start: int = cmtor.next_line()
+
+        # enter new iteration (mantain locals())
+        cmtor.append_source("")
+        cmtor.append_source("###### !new iteration! ######")
+        cmtor._locals[REG_it] = it
+        stmt = ast.Assign([self.target], ast.Name(REG_it, ast.Load()))
+        cmtor.process(stmt)
+
+        # process body
+        for stmt in self.body:
+            if type(stmt) not in flags.APPEND_SOURCE_BY_THEMSELVES:
+                cmtor.append_source(to_source(stmt))
+            if cmtor._stack_event == flags.NORMAL:
+                cmtor.process(stmt)
+            else:
+                cmtor._lines[-1] += "  # skipped"
+        cmtor.append_source()
+
+        if cmtor._stack_event == flags.BREAK:
+            cmtor._stack_event = flags.NORMAL
+            break
+        if cmtor._stack_event == flags.CONTINUE:
+            cmtor._stack_event = flags.NORMAL
+            continue
+
+    cmtor.indent -= flags.INDENT
+
+    # comment out all code except for the last iter
+    for lineno in range(loop_start, last_iter_start):
+        if cmtor._lines_category[lineno][0] == flags.SOURCE:
+            line: str = cmtor._lines[lineno]
+            if line.lstrip() and line.lstrip()[0] != "#":
+                cmtor._lines[lineno] = " " * self_indent + "# " + line[self_indent:]
+
+
+def Break(self, cmtor):
+    cmtor._lines[-1] += "  # True"
+    cmtor._stack_event = flags.BREAK
+
+
+def Continue(self, cmtor):
+    cmtor._lines[-1] += "  # True"
+    cmtor._stack_event = flags.CONTINUE

trace_commentor/handlers/definitions.py

@@ -1,6 +1,5 @@
-import ast
 from .. import flags
-from ..utils import to_source, APPEND_SOURCE_BY_THEMSELVES
+from ..utils import to_source
 
 def FunctionDef(self, cmtor):
     cmtor.append_source(f"def {self.name}():")
@@ -8,7 +7,7 @@ def FunctionDef(self, cmtor):
     
     for stmt in self.body:
         
-        if type(stmt) not in APPEND_SOURCE_BY_THEMSELVES:
+        if type(stmt) not in flags.APPEND_SOURCE_BY_THEMSELVES:
             cmtor.append_source(to_source(stmt))
         
         if self is cmtor.root:

trace_commentor/handlers/literals.py

@@ -5,3 +5,8 @@ def Constant(self, cmtor):
 def Tuple(self, cmtor):
     for x in self.elts:
         cmtor.process(x)
+
+
+def List(self, cmtor):
+    for x in self.elts:
+        cmtor.process(x)

trace_commentor/handlers/statements.py

@@ -1,5 +1,4 @@
-import ast
+from .. import flags
-from ..utils import to_source
 
 def Pass(self, cmtor):
     pass
@@ -7,7 +6,7 @@ def Pass(self, cmtor):
 def Assign(self, cmtor):
     cmtor.process(self.value)
     cmtor.exec(self)
-    if type(self.value) not in [ast.Constant]:
+    if type(self.value) not in flags.ASSIGN_SILENT:
-        cmtor.append_comment("----------")
+        cmtor.append_comment(f"----------")
         for target in self.targets:
             cmtor.process(target)

trace_commentor/utils.py

@@ -5,11 +5,6 @@ import os
 from . import flags
 
 
-APPEND_SOURCE_BY_THEMSELVES = [
-    ast.If,
-]
-
-
 def sign(line: str, depth=1):
     if flags.DEBUG:
         currentframe = inspect.currentframe()