Improve sql.util (#41)

Author: bryanforbes

sqlalchemy-stubs/sql/util.pyi

@@ -1,88 +1,107 @@
 from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import Iterator
+from typing import List
 from typing import Optional
+from typing import Sequence
+from typing import TypeVar
+from typing import Union
 
-from . import coercions as coercions
-from . import operators as operators
-from . import roles as roles
-from . import visitors as visitors
-from .base import ColumnSet as ColumnSet
-from .ddl import sort_tables as sort_tables
-from .elements import BindParameter as BindParameter
-from .elements import ColumnClause as ColumnClause
-from .elements import ColumnElement as ColumnElement
-from .elements import Grouping as Grouping
-from .elements import Label as Label
-from .elements import Null as Null
-from .elements import UnaryExpression as UnaryExpression
-from .schema import Column as Column
-from .selectable import Alias as Alias
-from .selectable import FromClause as FromClause
-from .selectable import FromGrouping as FromGrouping
-from .selectable import Join as Join
-from .selectable import ScalarSelect as ScalarSelect
-from .selectable import SelectBase as SelectBase
-from .selectable import TableClause as TableClause
-from .traversals import HasCacheKey as HasCacheKey
-from .. import exc as exc
-from .. import util as util
+from typing_extensions import Protocol
 
-join_condition: Any
+from . import visitors
+from .base import ColumnCollection
+from .base import ColumnSet
+from .elements import ClauseElement
+from .elements import ColumnClause
+from .schema import Table
+from .selectable import Alias
+from .selectable import FromClause
+from .selectable import Join
+from .selectable import Selectable
 
-def find_join_source(clauses: Any, join_to: Any): ...
-def find_left_clause_that_matches_given(clauses: Any, join_from: Any): ...
+_T = TypeVar("_T")
+_COLA = TypeVar("_COLA", bound=ColumnAdapter)
+
+class _IsDerivedFrom(Protocol):
+    def is_derived_from(self, __fromclause: FromClause) -> bool: ...
+
+join_condition = Join._join_condition
+
+def find_join_source(
+    clauses: Iterable[_IsDerivedFrom], join_to: Selectable
+) -> List[int]: ...
+def find_left_clause_that_matches_given(
+    clauses: Sequence[_IsDerivedFrom], join_from: Selectable
+) -> List[int]: ...
 def find_left_clause_to_join_from(
-    clauses: Any, join_to: Any, onclause: Any
-): ...
-def visit_binary_product(fn: Any, expr: Any) -> None: ...
+    clauses: Sequence[FromClause],
+    join_to: Selectable,
+    onclause: Optional[Any],
+) -> Iterable[int]: ...
+def visit_binary_product(
+    fn: Callable[[Any, Any, Any], Any], expr: ClauseElement
+) -> Iterator[ClauseElement]: ...
 def find_tables(
-    clause: Any,
+    clause: Union[ClauseElement, FromClause],
     check_columns: bool = ...,
     include_aliases: bool = ...,
     include_joins: bool = ...,
     include_selects: bool = ...,
     include_crud: bool = ...,
-): ...
-def unwrap_order_by(clause: Any): ...
-def unwrap_label_reference(element: Any): ...
-def expand_column_list_from_order_by(collist: Any, order_by: Any): ...
-def clause_is_present(clause: Any, search: Any): ...
-def tables_from_leftmost(clause: Any) -> None: ...
-def surface_selectables(clause: Any) -> None: ...
-def surface_selectables_only(clause: Any) -> None: ...
-def extract_first_column_annotation(column: Any, annotation_name: Any): ...
-def selectables_overlap(left: Any, right: Any): ...
-def bind_values(clause: Any): ...
+) -> List[Table]: ...
+def unwrap_order_by(clause: ClauseElement) -> List[ClauseElement]: ...
+def unwrap_label_reference(element: _T) -> _T: ...
+def expand_column_list_from_order_by(
+    collist: ColumnCollection[Any], order_by: ClauseElement
+) -> List[ClauseElement]: ...
+def clause_is_present(clause: Selectable, search: Selectable) -> bool: ...
+def tables_from_leftmost(clause: Table) -> Iterator[Table]: ...
+def surface_selectables(clause: Any) -> Iterator[FromClause]: ...
+def surface_selectables_only(
+    clause: Any,
+) -> Iterator[Union[Alias, FromClause, ColumnClause[Any]]]: ...
+def extract_first_column_annotation(
+    column: Any, annotation_name: str
+) -> Optional[Any]: ...
+def selectables_overlap(left: Any, right: Any) -> bool: ...
+def bind_values(clause: Any) -> List[Any]: ...
 
 class _repr_base:
-    def trunc(self, value: Any): ...
+    def trunc(self, value: Any) -> str: ...
 
 class _repr_row(_repr_base):
     row: Any = ...
-    max_chars: Any = ...
+    max_chars: int = ...
     def __init__(self, row: Any, max_chars: int = ...) -> None: ...
 
 class _repr_params(_repr_base):
     params: Any = ...
-    ismulti: Any = ...
+    ismulti: Optional[bool] = ...
     batches: Any = ...
-    max_chars: Any = ...
+    max_chars: int = ...
     def __init__(
         self,
         params: Any,
         batches: Any,
         max_chars: int = ...,
-        ismulti: Optional[Any] = ...,
+        ismulti: Optional[bool] = ...,
     ) -> None: ...
 
-def adapt_criterion_to_null(crit: Any, nulls: Any): ...
-def splice_joins(left: Any, right: Any, stop_on: Optional[Any] = ...): ...
-def reduce_columns(columns: Any, *clauses: Any, **kw: Any): ...
+def adapt_criterion_to_null(crit: Any, nulls: Any) -> Any: ...
+def splice_joins(
+    left: Any, right: Any, stop_on: Optional[Any] = ...
+) -> Any: ...
+def reduce_columns(
+    columns: Any, *clauses: Any, **kw: Any
+) -> ColumnSet[Any]: ...
 def criterion_as_pairs(
     expression: Any,
     consider_as_foreign_keys: Optional[Any] = ...,
     consider_as_referenced_keys: Optional[Any] = ...,
     any_operator: bool = ...,
-): ...
+) -> Any: ...
 
 class ClauseAdapter(visitors.ReplacingExternalTraversal):
     __traverse_options__: Any = ...
@@ -100,7 +119,7 @@ class ClauseAdapter(visitors.ReplacingExternalTraversal):
         adapt_on_names: bool = ...,
         anonymize_labels: bool = ...,
     ) -> None: ...
-    def replace(self, col: Any): ...
+    def replace(self, col: Any) -> Optional[Any]: ...
 
 class ColumnAdapter(ClauseAdapter):
     columns: Any = ...
@@ -121,9 +140,9 @@ class ColumnAdapter(ClauseAdapter):
         parent: Any = ...
         columns: Any = ...
         def __init__(self, parent: Any, columns: Any) -> None: ...
-        def __getitem__(self, key: Any): ...
-    def wrap(self, adapter: Any): ...
-    def traverse(self, obj: Any): ...
-    adapt_clause: Any = ...
-    adapt_list: Any = ...
-    def adapt_check_present(self, col: Any): ...
+        def __getitem__(self, key: Any) -> Any: ...
+    def wrap(self: _COLA, adapter: Any) -> _COLA: ...
+    def traverse(self, obj: Any) -> Any: ...
+    def adapt_clause(self, obj: Any) -> Any: ...
+    adapt_list = ClauseAdapter.copy_and_process
+    def adapt_check_present(self, col: Any) -> Optional[Any]: ...