Improve sql.lambdas (#35)

Author: bryanforbes

sqlalchemy-stubs/sql/lambdas.pyi

@@ -1,19 +1,13 @@
 from typing import Any
+from typing import Callable
+from typing import Iterator
 from typing import Optional
+from typing import Type
 
-from . import coercions as coercions
-from . import elements as elements
-from . import roles as roles
-from . import schema as schema
-from . import traversals as traversals
-from . import type_api as type_api
-from . import visitors as visitors
-from .base import Options as Options
-from .operators import ColumnOperators as ColumnOperators
-from .. import exc as exc
-from .. import inspection as inspection
-from .. import util as util
-from ..util import collections_abc as collections_abc
+from . import elements
+from . import roles
+from .base import Options
+from .operators import ColumnOperators
 
 class LambdaOptions(Options):
     enable_tracking: bool = ...
@@ -31,99 +25,123 @@ def lambda_stmt(
     global_track_bound_values: bool = ...,
     track_bound_values: bool = ...,
     lambda_cache: Optional[Any] = ...,
-): ...
+) -> StatementLambdaElement: ...
 
 class LambdaElement(elements.ClauseElement):
     __visit_name__: str = ...
     parent_lambda: Any = ...
-    fn: Any = ...
-    role: Any = ...
+    fn: Callable[..., elements.ClauseElement] = ...
+    role: Type[Any] = ...
     tracker_key: Any = ...
     opts: Any = ...
     def __init__(
         self,
-        fn: Any,
-        role: Any,
+        fn: Callable[..., elements.ClauseElement],
+        role: Type[Any],
         opts: Any = ...,
         apply_propagate_attrs: Optional[Any] = ...,
     ) -> None: ...
-    def __getattr__(self, key: Any): ...
+    def __getattr__(self, key: Any) -> Any: ...
 
 class DeferredLambdaElement(LambdaElement):
     lambda_args: Any = ...
     def __init__(
-        self, fn: Any, role: Any, opts: Any = ..., lambda_args: Any = ...
+        self,
+        fn: Callable[..., elements.ClauseElement],
+        role: Type[Any],
+        opts: Any = ...,
+        lambda_args: Any = ...,
     ) -> None: ...
 
 class StatementLambdaElement(roles.AllowsLambdaRole, LambdaElement):
-    def __add__(self, other: Any): ...
+    def __add__(self, other: Any) -> LinkedLambdaElement: ...
     def add_criteria(
         self,
         other: Any,
         enable_tracking: bool = ...,
         track_on: Optional[Any] = ...,
         track_closure_variables: bool = ...,
         track_bound_values: bool = ...,
-    ): ...
-    def spoil(self): ...
+    ) -> LinkedLambdaElement: ...
+    def spoil(self) -> NullLambdaStatement: ...
 
 class NullLambdaStatement(roles.AllowsLambdaRole, elements.ClauseElement):
     __visit_name__: str = ...
     def __init__(self, statement: Any) -> None: ...
-    def __getattr__(self, key: Any): ...
-    def __add__(self, other: Any): ...
-    def add_criteria(self, other: Any, **kw: Any): ...
+    def __getattr__(self, key: Any) -> Any: ...
+    def __add__(self, other: Any) -> NullLambdaStatement: ...
+    def add_criteria(self, other: Any, **kw: Any) -> NullLambdaStatement: ...
 
 class LinkedLambdaElement(StatementLambdaElement):
     role: Any = ...
     opts: Any = ...
     fn: Any = ...
     parent_lambda: Any = ...
     tracker_key: Any = ...
-    def __init__(self, fn: Any, parent_lambda: Any, opts: Any) -> None: ...
+    def __init__(
+        self,
+        fn: Callable[..., elements.ClauseElement],
+        parent_lambda: StatementLambdaElement,
+        opts: Any,
+    ) -> None: ...
 
 class AnalyzedCode:
     @classmethod
-    def get(cls, fn: Any, lambda_element: Any, lambda_kw: Any, **kw: Any): ...
+    def get(
+        cls,
+        fn: Callable[..., elements.ClauseElement],
+        lambda_element: LambdaElement,
+        lambda_kw: Any,
+        **kw: Any,
+    ) -> Any: ...
     track_bound_values: Any = ...
     track_closure_variables: Any = ...
     bindparam_trackers: Any = ...
     closure_trackers: Any = ...
     build_py_wrappers: Any = ...
-    def __init__(self, fn: Any, lambda_element: Any, opts: Any) -> None: ...
+    def __init__(
+        self,
+        fn: Callable[..., elements.ClauseElement],
+        lambda_element: LambdaElement,
+        opts: Any,
+    ) -> None: ...
 
 class AnalyzedFunction:
     analyzed_code: Any = ...
-    fn: Any = ...
+    fn: Callable[..., elements.ClauseElement] = ...
     bindparam_trackers: Any = ...
     def __init__(
         self,
         analyzed_code: Any,
-        lambda_element: Any,
+        lambda_element: LambdaElement,
         apply_propagate_attrs: Any,
-        fn: Any,
+        fn: Callable[..., elements.ClauseElement],
     ) -> None: ...
 
 class PyWrapper(ColumnOperators):
-    fn: Any = ...
+    fn: Callable[..., elements.ClauseElement] = ...
     track_bound_values: Any = ...
     def __init__(
         self,
-        fn: Any,
-        name: Any,
+        fn: Callable[..., elements.ClauseElement],
+        name: str,
         to_evaluate: Any,
         closure_index: Optional[Any] = ...,
         getter: Optional[Any] = ...,
         track_bound_values: bool = ...,
     ) -> None: ...
-    def __call__(self, *arg: Any, **kw: Any): ...
-    def operate(self, op: Any, *other: Any, **kwargs: Any): ...
-    def reverse_operate(self, op: Any, other: Any, **kwargs: Any): ...
-    def __clause_element__(self): ...
-    def __bool__(self): ...
-    def __nonzero__(self): ...
-    def __getattribute__(self, key: Any): ...
-    def __iter__(self) -> Any: ...
-    def __getitem__(self, key: Any): ...
+    def __call__(self, *arg: Any, **kw: Any) -> Any: ...
+    def operate(
+        self, op: Any, *other: Any, **kwargs: Any
+    ) -> elements.ClauseElement: ...
+    def reverse_operate(
+        self, op: Any, other: Any, **kwargs: Any
+    ) -> elements.ClauseElement: ...
+    def __clause_element__(self) -> elements.ClauseElement: ...
+    def __bool__(self) -> bool: ...
+    def __nonzero__(self) -> bool: ...
+    def __getattribute__(self, key: Any) -> Any: ...
+    def __iter__(self) -> Iterator[Any]: ...
+    def __getitem__(self, key: Any) -> Any: ...
 
-def insp(lmb: Any): ...
+def insp(lmb: Any) -> Any: ...