diff --git a/.gitignore b/.gitignore index 7e731c6..43070c2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ __pycache__ *.log +*.log.py .vscode-upload.json diff --git a/experiment.py b/experiment.py index 6d93872..d5eb021 100644 --- a/experiment.py +++ b/experiment.py @@ -5,20 +5,12 @@ def test(): @Commentor() 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 - + # return only odd numbers - 1,3,5,7,9 + for x in range(10): + # Check if x is even + if x % 2 == 0: + continue + print(x) print(target()) diff --git a/tests/test_conditions.py b/tests/test_conditions.py deleted file mode 100644 index 3ab78c1..0000000 --- a/tests/test_conditions.py +++ /dev/null @@ -1,61 +0,0 @@ -from test_utils import * - - -def test_constant(): - - @Commentor("") - 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("") - 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 -''') diff --git a/tests/test_control_flow.py b/tests/test_control_flow.py new file mode 100644 index 0000000..82cca0c --- /dev/null +++ b/tests/test_control_flow.py @@ -0,0 +1,208 @@ +from test_utils import * + + +def test_constant(): + + @Commentor("") + 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("") + 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 + """''') + + diff --git a/trace_commentor/commentor.py b/trace_commentor/commentor.py index 8cbb45c..32efac0 100644 --- a/trace_commentor/commentor.py +++ b/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): @@ -41,17 +44,25 @@ class Commentor(object): # input { self._locals = kwargs # } - + self.process(self.root) - + # output { code = "\n".join(self._lines) if self.file == "": return code elif self.file == "": - rich.print(code, file=sys.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 == "": - rich.print(code, file=sys.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) elif type(self.file) == str: @@ -76,12 +87,12 @@ class Commentor(object): obj = eval(src, self._globals, self._locals) if not format: return obj - + fmt = self.get_formatter(obj) fmt_obj = fmt(obj) if fmt_obj is not None: return f"{fmt(obj)} : {src}" - + def exec(self, node: ast.stmt): src = to_source(node) exec(src, self._globals, self._locals) @@ -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))) diff --git a/trace_commentor/flags.py b/trace_commentor/flags.py index 61b02b5..506a7f0 100644 --- a/trace_commentor/flags.py +++ b/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 +] diff --git a/trace_commentor/handlers/__init__.py b/trace_commentor/handlers/__init__.py index 3ce26e1..6612cfe 100644 --- a/trace_commentor/handlers/__init__.py +++ b/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 diff --git a/trace_commentor/handlers/control_flow.py b/trace_commentor/handlers/control_flow.py index ba6b902..696716c 100644 --- a/trace_commentor/handlers/control_flow.py +++ b/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 @@ -17,34 +16,95 @@ def If(self, cmtor, state=0): test_comment = test if test: state = state | PASS - + if state & ELIF: cmtor.append_source(f"elif {to_source(self.test)}: # {test_comment}") else: cmtor.append_source(f"if {to_source(self.test)}: # {test_comment}") - + 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 - + if self.orelse: if type(self.orelse[0]) == ast.If: cmtor.process(self.orelse[0], state=state | ELIF) else: test = not (state & PASS) test_comment = True if test else "skipped" - cmtor.append_source(f"else: # {test_comment}") - + cmtor.append_source(f"else: # {test_comment}") + 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 diff --git a/trace_commentor/handlers/definitions.py b/trace_commentor/handlers/definitions.py index 334e151..bcbd0aa 100644 --- a/trace_commentor/handlers/definitions.py +++ b/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: diff --git a/trace_commentor/handlers/literals.py b/trace_commentor/handlers/literals.py index a36c0df..b94e268 100644 --- a/trace_commentor/handlers/literals.py +++ b/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) diff --git a/trace_commentor/handlers/statements.py b/trace_commentor/handlers/statements.py index 649efdb..34bcc81 100644 --- a/trace_commentor/handlers/statements.py +++ b/trace_commentor/handlers/statements.py @@ -1,5 +1,4 @@ -import ast -from ..utils import to_source +from .. import flags 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]: - cmtor.append_comment("----------") + if type(self.value) not in flags.ASSIGN_SILENT: + cmtor.append_comment(f"----------") for target in self.targets: cmtor.process(target) diff --git a/trace_commentor/utils.py b/trace_commentor/utils.py index 27e77c4..6d31aa5 100644 --- a/trace_commentor/utils.py +++ b/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()