Improve jmespath (python#14176)

Author: donBarbos

stubs/jmespath/jmespath/__init__.pyi

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

stubs/jmespath/jmespath/ast.pyi

@@ -1,22 +1,55 @@
-def comparator(name, first, second): ...
-def current_node(): ...
-def expref(expression): ...
-def function_expression(name, args): ...
-def field(name): ...
-def filter_projection(left, right, comparator): ...
-def flatten(node): ...
-def identity(): ...
-def index(index): ...
-def index_expression(children): ...
-def key_val_pair(key_name, node): ...
-def literal(literal_value): ...
-def multi_select_dict(nodes): ...
-def multi_select_list(nodes): ...
-def or_expression(left, right): ...
-def and_expression(left, right): ...
-def not_expression(expr): ...
-def pipe(left, right): ...
-def projection(left, right): ...
-def subexpression(children): ...
-def slice(start, end, step): ...
-def value_projection(left, right): ...
+from typing import Literal, TypedDict
+from typing_extensions import NotRequired, TypeAlias
+
+_NodeType: TypeAlias = Literal[
+    "comparator",
+    "current",
+    "expref",
+    "function_expression",
+    "field",
+    "filter_projection",
+    "flatten",
+    "identity",
+    "index",
+    "index_expression",
+    "key_val_pair",
+    "literal",
+    "multi_select_dict",
+    "multi_select_list",
+    "or_expression",
+    "and_expression",
+    "not_expression",
+    "pipe",
+    "projection",
+    "subexpression",
+    "slice",
+    "value_projection",
+]
+
+class _ASTNode(TypedDict):
+    type: _NodeType
+    children: list[_ASTNode]
+    value: NotRequired[str]
+
+def comparator(name: str, first: _ASTNode, second: _ASTNode) -> _ASTNode: ...
+def current_node() -> _ASTNode: ...
+def expref(expression: _ASTNode) -> _ASTNode: ...
+def function_expression(name: str, args: list[_ASTNode]) -> _ASTNode: ...
+def field(name: str) -> _ASTNode: ...
+def filter_projection(left: _ASTNode, right: _ASTNode, comparator: _ASTNode) -> _ASTNode: ...
+def flatten(node: _ASTNode) -> _ASTNode: ...
+def identity() -> _ASTNode: ...
+def index(index: str) -> _ASTNode: ...
+def index_expression(children: list[_ASTNode]) -> _ASTNode: ...
+def key_val_pair(key_name: str, node: _ASTNode) -> _ASTNode: ...
+def literal(literal_value: str) -> _ASTNode: ...
+def multi_select_dict(nodes: list[_ASTNode]) -> _ASTNode: ...
+def multi_select_list(nodes: list[_ASTNode]) -> _ASTNode: ...
+def or_expression(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
+def and_expression(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
+def not_expression(expr: _ASTNode) -> _ASTNode: ...
+def pipe(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
+def projection(left: _ASTNode, right: _ASTNode) -> _ASTNode: ...
+def subexpression(children: list[_ASTNode]) -> _ASTNode: ...
+def slice(start: _ASTNode, end: _ASTNode, step: _ASTNode) -> _ASTNode: ...
+def value_projection(left, right) -> _ASTNode: ...

stubs/jmespath/jmespath/compat.pyi

@@ -1,10 +1,13 @@
-from _typeshed import Incomplete
 from collections.abc import Generator
 from itertools import zip_longest as zip_longest
+from types import FunctionType
+from typing import TypeVar
+
+_T = TypeVar("_T")
 
 text_type = str
 string_type = str
 
-def with_str_method(cls): ...
-def with_repr_method(cls): ...
-def get_methods(cls) -> Generator[Incomplete, None, None]: ...
+def with_str_method(cls: _T) -> _T: ...
+def with_repr_method(cls: _T) -> _T: ...
+def get_methods(cls: object) -> Generator[tuple[str, FunctionType]]: ...

stubs/jmespath/jmespath/functions.pyi

@@ -14,7 +14,7 @@ _F = TypeVar("_F", bound=Callable[..., Any])
 def signature(*arguments: _Signature) -> Callable[[_F], _F]: ...
 
 class FunctionRegistry(type):
-    def __init__(cls, name, bases, attrs) -> None: ...
+    def __init__(cls, name: str, bases: tuple[type, ...], attrs: dict[str, Any]) -> None: ...
 
 class Functions(metaclass=FunctionRegistry):
     FUNCTION_TABLE: Any

stubs/jmespath/jmespath/visitor.pyi

@@ -1,8 +1,11 @@
+from _typeshed import Unused
 from collections.abc import Callable, MutableMapping
-from typing import Any, ClassVar, NoReturn, TypedDict
+from typing import Any, ClassVar, NoReturn, TypedDict, TypeVar
 
 from jmespath.functions import Functions
 
+_T = TypeVar("_T")
+
 class Options:
     dict_cls: Callable[[], MutableMapping[Any, Any]] | None
     custom_functions: Functions | None
@@ -19,7 +22,7 @@ class _Expression:
 class Visitor:
     def __init__(self) -> None: ...
     def visit(self, node: _TreeNode, *args, **kwargs) -> Any: ...
-    def default_visit(self, node: _TreeNode, *args, **kwargs) -> NoReturn: ...
+    def default_visit(self, node: _TreeNode, *args: Unused, **kwargs: Unused) -> NoReturn: ...
 
 class _TreeNode(TypedDict):
     type: str
@@ -30,29 +33,29 @@ class TreeInterpreter(Visitor):
     COMPARATOR_FUNC: ClassVar[dict[str, Callable[[Any, Any], Any]]]
     MAP_TYPE: ClassVar[Callable[[], MutableMapping[Any, Any]]]
     def __init__(self, options: Options | None = None) -> None: ...
-    def default_visit(self, node: _TreeNode, *args, **kwargs) -> NoReturn: ...
+    def default_visit(self, node: _TreeNode, *args: Unused, **kwargs: Unused) -> NoReturn: ...
     def visit_subexpression(self, node: _TreeNode, value: Any) -> Any: ...
     def visit_field(self, node: _TreeNode, value: Any) -> Any: ...
     def visit_comparator(self, node: _TreeNode, value: Any) -> Any: ...
-    def visit_current(self, node: _TreeNode, value: Any) -> Any: ...
-    def visit_expref(self, node: _TreeNode, value: Any) -> Any: ...
+    def visit_current(self, node: _TreeNode, value: _T) -> _T: ...
+    def visit_expref(self, node: _TreeNode, value: Any) -> _Expression: ...
     def visit_function_expression(self, node: _TreeNode, value: Any) -> Any: ...
-    def visit_filter_projection(self, node: _TreeNode, value: Any) -> Any: ...
+    def visit_filter_projection(self, node: _TreeNode, value: Any) -> list[Any] | None: ...
     def visit_flatten(self, node: _TreeNode, value: Any) -> Any: ...
-    def visit_identity(self, node: _TreeNode, value: Any) -> Any: ...
+    def visit_identity(self, node: _TreeNode, value: _T) -> _T: ...
     def visit_index(self, node: _TreeNode, value: Any) -> Any: ...
     def visit_index_expression(self, node: _TreeNode, value: Any) -> Any: ...
     def visit_slice(self, node: _TreeNode, value: Any) -> Any: ...
     def visit_key_val_pair(self, node: _TreeNode, value: Any) -> Any: ...
     def visit_literal(self, node: _TreeNode, value: Any) -> Any: ...
     def visit_multi_select_dict(self, node: _TreeNode, value: Any) -> Any: ...
-    def visit_multi_select_list(self, node: _TreeNode, value: Any) -> Any: ...
+    def visit_multi_select_list(self, node: _TreeNode, value: Any) -> list[Any] | None: ...
     def visit_or_expression(self, node: _TreeNode, value: Any) -> Any: ...
     def visit_and_expression(self, node: _TreeNode, value: Any) -> Any: ...
-    def visit_not_expression(self, node: _TreeNode, value: Any) -> Any: ...
+    def visit_not_expression(self, node: _TreeNode, value: Any) -> bool: ...
     def visit_pipe(self, node: _TreeNode, value: Any) -> Any: ...
-    def visit_projection(self, node: _TreeNode, value: Any) -> Any: ...
-    def visit_value_projection(self, node: _TreeNode, value: Any) -> Any: ...
+    def visit_projection(self, node: _TreeNode, value: Any) -> list[Any] | None: ...
+    def visit_value_projection(self, node: _TreeNode, value: Any) -> list[Any] | None: ...
 
 class GraphvizVisitor(Visitor):
     def __init__(self) -> None: ...