Add types for jmespath internals (#9338)

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

stubs/jmespath/jmespath/__init__.pyi

@@ -1,7 +1,7 @@
 from typing import Any
 
 from jmespath import parser as parser
-from jmespath.visitor import Options as Options
+from jmespath.visitor import Options as Options, _TreeNode
 
 def compile(expression: str) -> parser.ParsedResult: ...
-def search(expression: str, data: Any, options: Any | None = ...) -> Any: ...
+def search(expression: str, data: Any, options: Options | None = ...) -> _TreeNode: ...

stubs/jmespath/jmespath/lexer.pyi

@@ -1,11 +1,19 @@
-from typing import Any
+from collections.abc import Iterator
+from typing import ClassVar
+from typing_extensions import TypedDict
 
 from jmespath.exceptions import EmptyExpressionError as EmptyExpressionError, LexerError as LexerError
 
+class _LexerTokenizeResult(TypedDict):
+    type: str
+    value: str
+    start: int
+    end: int
+
 class Lexer:
-    START_IDENTIFIER: Any
-    VALID_IDENTIFIER: Any
-    VALID_NUMBER: Any
-    WHITESPACE: Any
-    SIMPLE_TOKENS: Any
-    def tokenize(self, expression) -> None: ...
+    START_IDENTIFIER: ClassVar[set[str]]
+    VALID_IDENTIFIER: ClassVar[set[str]]
+    VALID_NUMBER: ClassVar[set[str]]
+    WHITESPACE: ClassVar[set[str]]
+    SIMPLE_TOKENS: ClassVar[dict[str, str]]
+    def tokenize(self, expression: str) -> Iterator[_LexerTokenizeResult]: ...

stubs/jmespath/jmespath/parser.pyi

@@ -1,15 +1,19 @@
-from typing import Any
+from collections.abc import Iterator
+from typing import Any, ClassVar
+
+from jmespath.lexer import _LexerTokenizeResult
+from jmespath.visitor import Options, _TreeNode
 
 class Parser:
-    BINDING_POWER: Any
-    tokenizer: Any
+    BINDING_POWER: ClassVar[dict[str, int]]
+    tokenizer: Iterator[_LexerTokenizeResult] | None
     def __init__(self, lookahead: int = ...) -> None: ...
-    def parse(self, expression): ...
+    def parse(self, expression: str) -> ParsedResult: ...
     @classmethod
     def purge(cls) -> None: ...
 
 class ParsedResult:
-    expression: Any
-    parsed: Any
-    def __init__(self, expression, parsed) -> None: ...
-    def search(self, value, options: Any | None = ...): ...
+    expression: str
+    parsed: _TreeNode
+    def __init__(self, expression: str, parsed: _TreeNode) -> None: ...
+    def search(self, value: Any, options: Options | None = ...) -> _TreeNode: ...

stubs/jmespath/jmespath/visitor.pyi

@@ -1,49 +1,60 @@
-from typing import Any
+from collections.abc import Callable, MutableMapping
+from typing import Any, ClassVar, NoReturn
+from typing_extensions import TypedDict
+
+from jmespath.functions import Functions
 
 class Options:
-    dict_cls: Any
-    custom_functions: Any
-    def __init__(self, dict_cls: Any | None = ..., custom_functions: Any | None = ...) -> None: ...
+    dict_cls: Callable[[], MutableMapping[Any, Any]] | None
+    custom_functions: Functions | None
+    def __init__(
+        self, dict_cls: Callable[[], MutableMapping[Any, Any]] | None = ..., custom_functions: Functions | None = ...
+    ) -> None: ...
 
 class _Expression:
-    expression: Any
-    interpreter: Any
-    def __init__(self, expression, interpreter) -> None: ...
-    def visit(self, node, *args, **kwargs): ...
+    expression: str
+    interpreter: Visitor
+    def __init__(self, expression: str, interpreter: Visitor) -> None: ...
+    def visit(self, node: _TreeNode, *args, **kwargs) -> _TreeNode: ...
 
 class Visitor:
     def __init__(self) -> None: ...
-    def visit(self, node, *args, **kwargs): ...
-    def default_visit(self, node, *args, **kwargs) -> None: ...
+    def visit(self, node: _TreeNode, *args, **kwargs) -> Any: ...
+    def default_visit(self, node: _TreeNode, *args, **kwargs) -> NoReturn: ...
+
+class _TreeNode(TypedDict):
+    type: str
+    value: Any
+    children: list[_TreeNode]
 
 class TreeInterpreter(Visitor):
-    COMPARATOR_FUNC: Any
-    MAP_TYPE: Any
-    def __init__(self, options: Any | None = ...) -> None: ...
-    def default_visit(self, node, *args, **kwargs) -> None: ...
-    def visit_subexpression(self, node, value): ...
-    def visit_field(self, node, value): ...
-    def visit_comparator(self, node, value): ...
-    def visit_current(self, node, value): ...
-    def visit_expref(self, node, value): ...
-    def visit_function_expression(self, node, value): ...
-    def visit_filter_projection(self, node, value): ...
-    def visit_flatten(self, node, value): ...
-    def visit_identity(self, node, value): ...
-    def visit_index(self, node, value): ...
-    def visit_index_expression(self, node, value): ...
-    def visit_slice(self, node, value): ...
-    def visit_key_val_pair(self, node, value): ...
-    def visit_literal(self, node, value): ...
-    def visit_multi_select_dict(self, node, value): ...
-    def visit_multi_select_list(self, node, value): ...
-    def visit_or_expression(self, node, value): ...
-    def visit_and_expression(self, node, value): ...
-    def visit_not_expression(self, node, value): ...
-    def visit_pipe(self, node, value): ...
-    def visit_projection(self, node, value): ...
-    def visit_value_projection(self, node, value): ...
+    COMPARATOR_FUNC: ClassVar[dict[str, Callable[[Any, Any], Any]]]
+    MAP_TYPE: ClassVar[Callable[[], MutableMapping[Any, Any]]]
+    def __init__(self, options: Options | None = ...) -> None: ...
+    def default_visit(self, node: _TreeNode, *args, **kwargs) -> NoReturn: ...
+    def visit_subexpression(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_field(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_comparator(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_current(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_expref(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_function_expression(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_filter_projection(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_flatten(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_identity(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_index(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_index_expression(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_slice(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_key_val_pair(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_literal(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_multi_select_dict(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_multi_select_list(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_or_expression(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_and_expression(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_not_expression(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_pipe(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_projection(self, node: _TreeNode, value: Any) -> _TreeNode: ...
+    def visit_value_projection(self, node: _TreeNode, value: Any) -> _TreeNode: ...
 
 class GraphvizVisitor(Visitor):
     def __init__(self) -> None: ...
-    def visit(self, node, *args, **kwargs): ...
+    def visit(self, node: _TreeNode, *args, **kwargs) -> str: ...