feat: improve rule typing

part of #389
ast-grep · Oct 23, 2023 · 80b255a · 80b255a
1 parent 097a209
commit 80b255a
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/crates/pyo3/ast_grep_pyo3.pyi b/crates/pyo3/ast_grep_pyo3.pyi
@@ -1,5 +1,12 @@
-from typing import Optional
+from typing import Optional, TypedDict, Unpack, NotRequired
 
+class Pattern(TypedDict):
+  selector: str
+  context: str
+
+class Rule(TypedDict):
+  pattern: NotRequired[str | Pattern]
+  kind: NotRequired[str]
 
 class Pos:
   line: int
@@ -16,6 +23,6 @@ class SgRoot:
 
 class SgNode:
   def range(self) -> Range: pass
-  def find(self, config = None, **kwargs) -> SgNode: ...
+  def find(self, config = None, **kwargs: Unpack[Rule]) -> SgNode: ...
   def get_match(self, meta_var: str) -> Optional[SgNode]: ...
   def text(self) -> str: ...