Add source field to ParserVisitor

openrewrite · Aug 12, 2024 · 163901b · 163901b
1 parent 9023939
commit 163901b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/rewrite/rewrite-python/src/rewrite/python/parser/parser_visitor.py b/rewrite/rewrite-python/src/rewrite/python/parser/parser_visitor.py
@@ -1,4 +1,5 @@
 import ast
+from dataclasses import dataclass
 from pathlib import Path
 
 import rewrite.java.tree as J
@@ -13,6 +14,13 @@ def map_type(node):
 
 
 class ParserVisitor(ast.NodeVisitor):
+    _source: str
+    _cursor: int = 0
+
+    def __init__(self, source: str):
+        super().__init__()
+        self._source = source
+
     def visit_Module(self, node: ast.Module) -> Py.CompilationUnit:
         return Py.CompilationUnit(
             random_id(),

diff --git a/rewrite/rewrite-python/tests/parser_test.py b/rewrite/rewrite-python/tests/parser_test.py
@@ -15,5 +15,5 @@ def bar(x):
     tree = ast.parse(source)
 
     # Create the visitor and visit the AST
-    visitor = ParserVisitor()
+    visitor = ParserVisitor(source)
     visitor.visit(tree)