Fix typing for ColumnCollection and subclasses (#65)

Author: bryanforbes

sqlalchemy-stubs/sql/base.pyi

@@ -16,7 +16,7 @@ from typing import TypeVar
 from typing import Union
 
 from . import roles
-from .schema import Column
+from .elements import ColumnElement
 from .traversals import HasCacheKey
 from .traversals import HasCopyInternals
 from .visitors import ClauseVisitor
@@ -30,8 +30,8 @@ _T = TypeVar("_T")
 _SC = TypeVar("_SC", bound=SingletonConstant)
 _O = TypeVar("_O", bound=Options)
 _E = TypeVar("_E", bound=Executable)
-_C = TypeVar("_C", bound=Column[Any])
-_OC = TypeVar("_OC", bound=Column[Any])
+_CE = TypeVar("_CE", bound=ColumnElement[Any])
+_OCE = TypeVar("_OCE", bound=ColumnElement[Any])
 
 coercions: ModuleType
 elements: ModuleType
@@ -145,56 +145,58 @@ class SchemaEventTarget: ...
 class SchemaVisitor(ClauseVisitor):
     __traverse_options__: Any = ...
 
-class ColumnCollection(Generic[_C]):
+class ColumnCollection(Generic[_CE]):
     def __init__(
-        self, columns: Optional[Iterable[Tuple[str, _C]]] = ...
+        self, columns: Optional[Iterable[Tuple[str, _CE]]] = ...
     ) -> None: ...
-    def keys(self) -> List[_C]: ...
+    def keys(self) -> List[_CE]: ...
     def values(self) -> List[str]: ...
-    def items(self) -> List[Tuple[str, _C]]: ...
+    def items(self) -> List[Tuple[str, _CE]]: ...
     def __bool__(self) -> bool: ...
     def __len__(self) -> int: ...
-    def __iter__(self) -> Iterator[_C]: ...
-    def __getitem__(self, key: str) -> _C: ...
-    def __getattr__(self, key: str) -> _C: ...
+    def __iter__(self) -> Iterator[_CE]: ...
+    def __getitem__(self, key: str) -> _CE: ...
+    def __getattr__(self, key: str) -> _CE: ...
     def __contains__(self, key: str) -> bool: ...
     def compare(self, other: Any) -> bool: ...
     def __eq__(self, other: Any) -> bool: ...
     @overload
-    def get(self, key: str) -> Optional[_C]: ...
+    def get(self, key: str) -> Optional[_CE]: ...
     @overload
-    def get(self, key: str, default: _T) -> Union[_C, _T]: ...
-    def __setitem__(self, key: str, value: _C) -> None: ...
+    def get(self, key: str, default: _T) -> Union[_CE, _T]: ...
+    def __setitem__(self, key: str, value: _CE) -> None: ...
     def __delitem__(self, key: str) -> None: ...
-    def __setattr__(self, key: str, obj: _C) -> None: ...
+    def __setattr__(self, key: str, obj: _CE) -> None: ...
     def clear(self) -> None: ...
-    def remove(self, column: _C) -> None: ...
+    def remove(self, column: _CE) -> None: ...
     def update(
-        self, iter_: Union[Mapping[str, _C], Iterable[Tuple[str, _C]]]
+        self, iter_: Union[Mapping[str, _CE], Iterable[Tuple[str, _CE]]]
     ) -> None: ...
     __hash__: Any = ...
-    def add(self, column: _C, key: Optional[str] = ...) -> None: ...
-    def contains_column(self, col: Column[Any]) -> bool: ...
-    def as_immutable(self) -> ImmutableColumnCollection[_C]: ...
+    def add(self, column: _CE, key: Optional[str] = ...) -> None: ...
+    def contains_column(self, col: ColumnElement[Any]) -> bool: ...
+    def as_immutable(self) -> ImmutableColumnCollection[_CE]: ...
     def corresponding_column(
-        self, column: Column[Any], require_embedded: bool = ...
-    ) -> Optional[_C]: ...
-
-class DedupeColumnCollection(ColumnCollection[_C]):
-    def add(self, column: _C, key: Optional[str] = ...) -> None: ...
-    def extend(self, iter_: Iterable[_C]) -> None: ...
-    def remove(self, column: _C) -> None: ...
-    def replace(self, column: _C) -> None: ...
-
-class ImmutableColumnCollection(util.ImmutableContainer, ColumnCollection[_C]):
-    def __init__(self, collection: ColumnCollection[_C]) -> None: ...
-    def add(self, column: _C, key: Optional[str] = ...) -> NoReturn: ...
-    def extend(self, iter_: Iterable[_C]) -> NoReturn: ...
-    def remove(self, column: _C) -> NoReturn: ...
-
-class ColumnSet(util.ordered_column_set[_C]):
-    def contains_column(self, col: Column[Any]) -> bool: ...
-    def extend(self, cols: Iterable[_C]) -> None: ...
-    def __add__(self, other: Iterable[_OC]) -> List[Union[_C, _OC]]: ...  # type: ignore[override]
+        self, column: ColumnElement[Any], require_embedded: bool = ...
+    ) -> Optional[_CE]: ...
+
+class DedupeColumnCollection(ColumnCollection[_CE]):
+    def add(self, column: _CE, key: Optional[str] = ...) -> None: ...
+    def extend(self, iter_: Iterable[_CE]) -> None: ...
+    def remove(self, column: _CE) -> None: ...
+    def replace(self, column: _CE) -> None: ...
+
+class ImmutableColumnCollection(
+    util.ImmutableContainer, ColumnCollection[_CE]
+):
+    def __init__(self, collection: ColumnCollection[_CE]) -> None: ...
+    def add(self, column: _CE, key: Optional[str] = ...) -> NoReturn: ...
+    def extend(self, iter_: Iterable[_CE]) -> NoReturn: ...
+    def remove(self, column: _CE) -> NoReturn: ...
+
+class ColumnSet(util.ordered_column_set[_CE]):
+    def contains_column(self, col: ColumnElement[Any]) -> bool: ...
+    def extend(self, cols: Iterable[_CE]) -> None: ...
+    def __add__(self, other: Iterable[_OCE]) -> List[Union[_CE, _OCE]]: ...  # type: ignore[override]
     def __eq__(self, other: Any) -> Any: ...
     def __hash__(self) -> int: ...  # type: ignore[override]