Also return comments from pythonparse.parse_buffer.

whitequark · whitequark · commit d63349cc6622 · 2015-06-06T15:27:05.000+03:00
diff --git a/pythonparser/__init__.py b/pythonparser/__init__.py
@@ -4,7 +4,12 @@
 def parse_buffer(buffer, mode="exec", flags=[], version=None, engine=None):
     """
     Like :meth:`parse`, but accepts a :class:`source.Buffer` instead of
-    source and filename."""
+    source and filename, and returns comments as well.
+
+    :see: :meth:`parse`
+    :return: (:class:`ast.AST`, list of :class:`source.Comment`)
+        Abstract syntax tree and comments
+    """
 
     if version is None:
         version = sys.version_info[0:2]
@@ -20,11 +25,11 @@ def parse_buffer(buffer, mode="exec", flags=[], version=None, engine=None):
     parser.add_flags(flags)
 
     if mode == "exec":
-        return parser.file_input()
+        return parser.file_input(), lexer.comments
     elif mode == "single":
-        return parser.single_input()
+        return parser.single_input(), lexer.comments
     elif mode == "eval":
-        return parser.eval_input()
+        return parser.eval_input(), lexer.comments
 
 def parse(source, filename="<unknown>", mode="exec",
           flags=[], version=None, engine=None):
@@ -49,6 +54,7 @@ def parse(source, filename="<unknown>", mode="exec",
     :raise: :class:`diagnostic.Error`
         if the source code is not well-formed
     """
-    return parse_buffer(pythonparser.source.Buffer(source, filename),
-                        mode, flags, version, engine)
+    ast, comments = parse_buffer(pythonparser.source.Buffer(source, filename),
+                                 mode, flags, version, engine)
+    return ast
 
diff --git a/pythonparser/lexer.py b/pythonparser/lexer.py
@@ -310,7 +310,7 @@ def _refill(self, eof_token):
             return
 
         if match.group(4) is not None: # comment
-            self.comments.append((tok_range, match.group(4)))
+            self.comments.append(source.Comment(tok_range, match.group(4)))
             return self._refill(eof_token)
 
         # Lexing non-whitespace now.
diff --git a/pythonparser/source.py b/pythonparser/source.py
@@ -170,6 +170,17 @@ def __ne__(self, other):
         """
         return not (self == other)
 
+class Comment:
+    """
+    A comment in the source code.
+
+    :ivar loc: (:class:`Range`) source location
+    :ivar text: (string) comment text
+    """
+
+    def __init__(self, loc, text):
+        self.loc, self.text = loc, text
+
 class RewriterConflict(Exception):
     """
     An exception that is raised when two ranges supplied to a rewriter overlap.
diff --git a/pythonparser/test/test_lexer.py b/pythonparser/test/test_lexer.py
@@ -63,8 +63,10 @@ def test_newline(self):
 
     def test_comment(self):
         self.assertLexes("# foo")
-        self.assertEqual([(source.Range(self.buffer, 0, 5), "# foo")],
-                         self.lexer.comments)
+        self.assertEqual(source.Range(self.buffer, 0, 5),
+                         self.lexer.comments[0].loc)
+        self.assertEqual("# foo",
+                         self.lexer.comments[0].text)
 
         self.assertLexes("class x:\n  # foo\n  pass",
                          "class",   None,
