feat: support get root

part of #389
ast-grep · Oct 23, 2023 · f1cf76e · f1cf76e
1 parent 40028b0
commit f1cf76e
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 3 deletions.
diff --git a/crates/pyo3/ast_grep_pyo3.pyi b/crates/pyo3/ast_grep_pyo3.pyi
@@ -17,4 +17,5 @@ class SgRoot:
 class SgNode:
   def range(self) -> Range: pass
   def find(self, config = None, **kwargs) -> SgNode: ...
-  def get_match(self, meta_var: str) -> Optional[SgNode]: ...
+  def get_match(self, meta_var: str) -> Optional[SgNode]: ...
+  def text(self) -> str: ...
diff --git a/crates/pyo3/src/py_node.rs b/crates/pyo3/src/py_node.rs
@@ -94,7 +94,9 @@ impl SgNode {
   }
 
   /*---------- Tree Traversal  ----------*/
-  // TODO get_root
+  fn get_root(&self) -> Py<SgRoot> {
+    self.root.clone()
+  }
 
   #[pyo3(signature = (config=None, **kwargs))]
   fn find(&self, config: Option<&PyDict>, kwargs: Option<&PyDict>) -> Option<Self> {

diff --git a/crates/pyo3/tests/test_simple.py b/crates/pyo3/tests/test_simple.py
@@ -24,4 +24,14 @@ def test_get_match():
   assert a.text() == 'a'
   rng = a.range()
   assert rng.start.line == 1
-  assert rng.start.column == 6
+  assert rng.start.column == 6
+
+def test_get_multi_match():
+  pass
+
+def test_get_root():
+  node = root.find(pattern = 'let a = $A')
+  assert node is not None
+  root2 = node.get_root()
+  assert root2.filename() == 'anonymous'
+  # assert root2 == root