Improve typing for sql.visitors

Author: bryanforbes

sqlalchemy-stubs/sql/visitors.pyi

@@ -1,92 +1,110 @@
 from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Iterable
+from typing import Iterator
+from typing import List
+from typing import Optional
+from typing import TypeVar
+
+_ET = TypeVar("_ET", bound=ExternalTraversal)
+_T = TypeVar("_T")
 
 class TraversibleType(type):
     def __init__(cls, clsname: Any, bases: Any, clsdict: Any) -> None: ...
 
-class Traversible: ...
+class Traversible(metaclass=TraversibleType): ...
 
 class _InternalTraversalType(type):
     def __init__(cls, clsname: Any, bases: Any, clsdict: Any) -> None: ...
 
-class InternalTraversal:
-    def dispatch(self, visit_symbol: Any): ...
+class InternalTraversal(object, metaclass=_InternalTraversalType):
+    def dispatch(self, visit_symbol: Any) -> Optional[Callable[..., Any]]: ...
     def run_generated_dispatch(
         self,
         target: Any,
         internal_dispatch: Any,
         generate_dispatcher_name: Any,
-    ): ...
+    ) -> Any: ...
     def generate_dispatch(
         self,
         target_cls: Any,
         internal_dispatch: Any,
         generate_dispatcher_name: Any,
-    ): ...
-    dp_has_cache_key: Any = ...
-    dp_has_cache_key_list: Any = ...
-    dp_clauseelement: Any = ...
-    dp_fromclause_canonical_column_collection: Any = ...
-    dp_clauseelement_tuples: Any = ...
-    dp_clauseelement_list: Any = ...
-    dp_clauseelement_tuple: Any = ...
-    dp_executable_options: Any = ...
-    dp_fromclause_ordered_set: Any = ...
-    dp_string: Any = ...
-    dp_string_list: Any = ...
-    dp_anon_name: Any = ...
-    dp_boolean: Any = ...
-    dp_operator: Any = ...
-    dp_type: Any = ...
-    dp_plain_dict: Any = ...
-    dp_dialect_options: Any = ...
-    dp_string_clauseelement_dict: Any = ...
-    dp_string_multi_dict: Any = ...
-    dp_annotations_key: Any = ...
-    dp_plain_obj: Any = ...
-    dp_named_ddl_element: Any = ...
-    dp_prefix_sequence: Any = ...
-    dp_table_hint_list: Any = ...
-    dp_setup_join_tuple: Any = ...
-    dp_statement_hint_list: Any = ...
-    dp_unknown_structure: Any = ...
-    dp_dml_ordered_values: Any = ...
-    dp_dml_values: Any = ...
-    dp_dml_multi_values: Any = ...
-    dp_propagate_attrs: Any = ...
+    ) -> Any: ...
+    dp_has_cache_key: int = ...
+    dp_has_cache_key_list: int = ...
+    dp_clauseelement: int = ...
+    dp_fromclause_canonical_column_collection: int = ...
+    dp_clauseelement_tuples: int = ...
+    dp_clauseelement_list: int = ...
+    dp_clauseelement_tuple: int = ...
+    dp_executable_options: int = ...
+    dp_fromclause_ordered_set: int = ...
+    dp_string: int = ...
+    dp_string_list: int = ...
+    dp_anon_name: int = ...
+    dp_boolean: int = ...
+    dp_operator: int = ...
+    dp_type: int = ...
+    dp_plain_dict: int = ...
+    dp_dialect_options: int = ...
+    dp_string_clauseelement_dict: int = ...
+    dp_string_multi_dict: int = ...
+    dp_annotations_key: int = ...
+    dp_plain_obj: int = ...
+    dp_named_ddl_element: int = ...
+    dp_prefix_sequence: int = ...
+    dp_table_hint_list: int = ...
+    dp_setup_join_tuple: int = ...
+    dp_statement_hint_list: int = ...
+    dp_unknown_structure: int = ...
+    dp_dml_ordered_values: int = ...
+    dp_dml_values: int = ...
+    dp_dml_multi_values: int = ...
+    dp_propagate_attrs: int = ...
 
 class ExtendedInternalTraversal(InternalTraversal):
-    dp_ignore: Any = ...
-    dp_inspectable: Any = ...
-    dp_multi: Any = ...
-    dp_multi_list: Any = ...
-    dp_has_cache_key_tuples: Any = ...
-    dp_inspectable_list: Any = ...
+    dp_ignore: int = ...
+    dp_inspectable: int = ...
+    dp_multi: int = ...
+    dp_multi_list: int = ...
+    dp_has_cache_key_tuples: int = ...
+    dp_inspectable_list: int = ...
 
 class ExternalTraversal:
     __traverse_options__: Any = ...
-    def traverse_single(self, obj: Any, **kw: Any): ...
-    def iterate(self, obj: Any): ...
-    def traverse(self, obj: Any): ...
+    def traverse_single(self, obj: Any, **kw: Any) -> Any: ...
+    def iterate(self, obj: Any) -> Iterator[Any]: ...
+    def traverse(self, obj: _T) -> _T: ...
     @property
-    def visitor_iterator(self) -> None: ...
-    def chain(self, visitor: Any): ...
+    def visitor_iterator(self) -> Iterator[Any]: ...
+    def chain(self: _ET, visitor: Any) -> _ET: ...
 
 class CloningExternalTraversal(ExternalTraversal):
-    def copy_and_process(self, list_: Any): ...
-    def traverse(self, obj: Any): ...
+    def copy_and_process(self, list_: Iterable[_T]) -> List[_T]: ...
+    def traverse(self, obj: _T) -> _T: ...
 
 class ReplacingExternalTraversal(CloningExternalTraversal):
-    def replace(self, elem: Any) -> None: ...
-    def traverse(self, obj: Any): ...
+    def replace(self, elem: Any) -> Optional[Any]: ...
+    def traverse(self, obj: _T) -> _T: ...
 
 Visitable = Traversible
 VisitableType = TraversibleType
 ClauseVisitor = ExternalTraversal
 CloningVisitor = CloningExternalTraversal
 ReplacingCloningVisitor = ReplacingExternalTraversal
 
-def iterate(obj: Any, opts: Any = ...) -> None: ...
-def traverse_using(iterator: Any, obj: Any, visitors: Any): ...
-def traverse(obj: Any, opts: Any, visitors: Any): ...
-def cloned_traverse(obj: Any, opts: Any, visitors: Any): ...
-def replacement_traverse(obj: Any, opts: Any, replace: Any): ...
+def iterate(obj: Any, opts: Dict[str, Any] = ...) -> Iterator[Any]: ...
+def traverse_using(
+    iterator: Iterator[Any], obj: _T, visitors: Dict[str, Callable[..., Any]]
+) -> _T: ...
+def traverse(
+    obj: _T, opts: Dict[str, Any], visitors: Dict[str, Callable[..., Any]]
+) -> _T: ...
+def cloned_traverse(
+    obj: _T, opts: Dict[str, Any], visitors: Dict[str, Callable[..., Any]]
+) -> _T: ...
+def replacement_traverse(
+    obj: _T, opts: Dict[str, Any], replace: Callable[[Any], Any]
+) -> _T: ...