Less hacky way of getting comments

tools/code_processing/loader.py

@@ -1,53 +1,18 @@
-import re
 import ast
 
 from . import python_to_ts, pseudocode
 
-RE_COMMENT = re.compile(r"(\s*)# ?(.*)")
-RE_NOLINE = re.compile(r"^(\s*)$", flags=re.MULTILINE)
-
-
-def indent_at(string, pos):
-    if pos <= 0:
-        return ""
-
-    newline_before = string.rfind('\n', 0, pos - 1)
-    if newline_before == -1:
-        return ""
-
-    line_before = string[newline_before + 1:pos]
-
-    additional = 0
-    if line_before.endswith(":\n"):
-        additional = 4
-
-    return " " * (len(line_before) - len(line_before.lstrip()) + additional)
-
-
-def gather_indent(m):
-    spaces = m.group(1)
-    if spaces:
-        return spaces
-
-    return indent_at(m.string, m.start())
-
-
-def process_code(source):
-    return RE_NOLINE.sub(
-        lambda m: gather_indent(m) + "''",
-        RE_COMMENT.sub(r"\1'''\2'''", source.strip())
-    )
-
 
 def code_to_ast(source):
-    return ast.parse(process_code(source), type_comments=True)
+    return ast.parse(source, type_comments=True)
 
 
 def code_to_samples(source):
     tree = code_to_ast(source)
+    comments = python_to_ts.CommentData(source)
     return {
         "ast": tree,
-        "pseudo": pseudocode.PseudoCode().convert(tree),
+        "pseudo": pseudocode.PseudoCode().convert(tree, comments),
         "py": source.strip("\n"),
-        "ts": python_to_ts.Py2Ts().convert(tree),
+        "ts": python_to_ts.Py2Ts().convert(tree, comments),
     }

tools/code_processing/pseudocode.py

@@ -214,3 +214,6 @@ def other_expression(self, value, annotation):
     def begin_for(self, target, iter, is_async):
         code_start = "For each $%s$ in $%s$" % (target, iter)
         self.push_code(code_start)
+
+    def convert_line_comment(self, comment):
+        return comment or ""

tools/code_processing/python_to_ts.py

@@ -27,6 +27,78 @@ def __exit__(self, *a):
         self.converter.indent_down()
 
 
+class CommentData:
+    RE_COMMENT = re.compile(r"^\s*(.+)?# ?(.*)$")
+
+    class Comment:
+        def __init__(self, line, text, is_tail):
+            self.line = line
+            self.text = text
+            self.is_tail = is_tail
+
+    class CommentReader:
+        def __init__(self, comments):
+            self.comments = comments
+            self.next_line = None
+            self.current_line = 0
+            self.index = 0
+            self.get_next()
+
+        def prepare(self, node):
+            if hasattr(node, "lineno"):
+                self.current_line = node.lineno
+
+        def get_next(self):
+            self.next_line = self.current_comment().line if self.index < len(self.comments) else None
+
+        def current_comment(self):
+            return self.comments[self.index]
+
+        def has_line_comment(self):
+            if self.next_line is None:
+                return False
+
+            if self.current_line >= self.next_line:
+                return not self.current_comment().is_tail
+
+            return False
+
+        def has_tail_comment(self):
+            if self.next_line is None:
+                return False
+
+            if self.current_line >= self.next_line:
+                return self.current_comment().is_tail
+
+            return False
+
+        def get_comment(self):
+            return self.current_comment().text
+
+        def next(self):
+            self.index += 1
+            self.get_next()
+
+    def __init__(self, source):
+        self.comments = []
+
+        for lineno, line in enumerate(source.splitlines()):
+            match = self.RE_COMMENT.match(line)
+            if match:
+                self.add_comment(lineno + 1, match.group(2), match.group(1) is not None)
+            elif line.strip() == "":
+                self.add_empty(lineno + 1)
+
+    def add_comment(self, line, text, is_tail):
+        self.comments.append(self.Comment(line, text, is_tail))
+
+    def add_empty(self, line):
+        self.comments.append(self.Comment(line, None, False))
+
+    def reader(self):
+        return self.CommentReader(self.comments)
+
+
 class Range:
     def __init__(self, ast_node, translator):
         self.ast_node = ast_node
@@ -61,8 +133,14 @@ def __init__(self):
         self.in_method = False
         self.scope = [{}]
 
-    def convert(self, obj):
+    def convert(self, obj, comments: CommentData):
         self.output = ""
+        self.comments = comments.reader()
+        self.next_comment = None
+        if comments.comments:
+            self.next_comment = comments.comments[0].line
+
+        self.comment_line = 0
         self.convert_ast(obj)
         return self.output
 
@@ -102,8 +180,22 @@ def unknown(self, obj, as_string=False):
         #     return unknown
         # self.ts_code += unknown
 
-    def push_code(self, ts_code):
-        self.output += self.indent + ts_code + "\n"
+    def push_code(self, code):
+        while self.comments.has_line_comment():
+            comment = self.comments.get_comment()
+            if comment is None:
+                self.output += "\n"
+            else:
+                self.output += self.indent + self.convert_line_comment(comment) + "\n"
+            self.comments.next()
+
+        self.output += self.indent + code
+
+        if self.comments.has_tail_comment():
+            self.output += " " + self.convert_line_comment(self.comments.get_comment())
+            self.comments.next()
+
+        self.output += "\n"
 
     def indent_up(self):
         self.indent_level += 1
@@ -195,6 +287,8 @@ def expression_statement(self, v):
         self.push_code(v)
 
     def convert_ast(self, obj):
+        self.comments.prepare(obj)
+
         if isinstance(obj, ast.Module):
             self.convert_ast(obj.body)
         elif isinstance(obj, list):
@@ -315,6 +409,9 @@ def other_expression(self, value, annotation):
     def expr_attribute(self, object, member):
         return "%s.%s" % (object, member)
 
+    def convert_line_comment(self, comment):
+        raise NotImplementedError
+
     def expression_to_string(self, value, annotation=False):
         if isinstance(value, ast.Constant):
             return self.convert_constant(value.value)
@@ -593,3 +690,6 @@ def format_comment(self, value):
 
     def expression_statement(self, v):
         self.push_code(v + ";")
+
+    def convert_line_comment(self, comment):
+        return "// " + comment