Improve typing for sql.type_api classes

bryanforbes · bryanforbes · commit b515784f3924 · 2021-03-24T09:26:00.000-05:00
diff --git a/sqlalchemy-stubs/sql/sqltypes.pyi b/sqlalchemy-stubs/sql/sqltypes.pyi
@@ -7,6 +7,8 @@ from typing import Any
 from typing import List
 from typing import Mapping
 from typing import Optional
+from typing import Type
+from typing import TypeVar
 from typing import Union
 
 from . import coercions as coercions
@@ -35,6 +37,8 @@ from ..util import compat as compat
 from ..util import langhelpers as langhelpers
 from ..util import pickle as pickle
 
+_U = TypeVar("_U")
+
 class _LookupExpressionAdapter:
     class Comparator(TypeEngine.Comparator): ...
     comparator_factory: Any = ...
@@ -160,7 +164,6 @@ class _Binary(TypeEngine[bytes]):
     def python_type(self): ...
     def bind_processor(self, dialect: Any): ...
     def result_processor(self, dialect: Any, coltype: Any): ...
-    def result_processor(self, dialect: Any, coltype: Any): ...
     def coerce_compared_value(self, op: Any, value: Any): ...
     def get_dbapi_type(self, dbapi: Any): ...
 
@@ -183,7 +186,7 @@ class SchemaType(SchemaEventTarget):
         _create_events: bool = ...,
     ) -> None: ...
     def copy(self, **kw: Any): ...
-    def adapt(self, impltype: Any, **kw: Any): ...
+    def adapt(self, __impltype: Type[_U], **kw: Any) -> _U: ...
     @property
     def bind(self): ...
     def create(
@@ -204,7 +207,7 @@ class Enum(Emulated, String, SchemaType):
     comparator_factory: Any = ...
     def as_generic(self, allow_nulltype: bool = ...): ...
     def adapt_to_emulated(self, impltype: Any, **kw: Any): ...
-    def adapt(self, impltype: Any, **kw: Any): ...
+    def adapt(self, __impltype: Type[_U], **kw: Any) -> _U: ...
     def literal_processor(self, dialect: Any): ...
     def bind_processor(self, dialect: Any): ...
     def result_processor(self, dialect: Any, coltype: Any): ...
@@ -248,20 +251,21 @@ class Boolean(Emulated, TypeEngine[bool], SchemaType):
 class _AbstractInterval(_LookupExpressionAdapter, TypeEngine[Any]):
     def coerce_compared_value(self, op: Any, value: Any): ...
 
-class Interval(Emulated, _AbstractInterval, TypeDecorator[timedelta]):
+# "comparator_factory" of "_LookupExpressionAdapter" and "TypeDecorator" are incompatible
+class Interval(Emulated, _AbstractInterval, TypeDecorator[timedelta]):  # type: ignore[misc]
     impl: Any = ...
-    epoch: Any = ...
-    native: Any = ...
-    second_precision: Any = ...
-    day_precision: Any = ...
+    epoch: datetime = ...
+    native: bool = ...
+    second_precision: Optional[float] = ...
+    day_precision: Optional[float] = ...
     def __init__(
         self,
         native: bool = ...,
-        second_precision: Optional[Any] = ...,
-        day_precision: Optional[Any] = ...,
+        second_precision: Optional[float] = ...,
+        day_precision: Optional[float] = ...,
     ) -> None: ...
     @property
-    def python_type(self): ...
+    def python_type(self) -> Type[timedelta]: ...
     def adapt_to_emulated(self, impltype: Any, **kw: Any): ...
     def bind_processor(self, dialect: Any): ...
     def result_processor(self, dialect: Any, coltype: Any): ...
@@ -270,7 +274,8 @@ class JSON(Indexable, TypeEngine[Union[str, Mapping, List]]):
     __visit_name__: str = ...
     hashable: bool = ...
     NULL: Any = ...
-    none_as_null: Any = ...
+    none_as_null: bool = ...
+    should_evaluate_none: bool = ...
     def __init__(self, none_as_null: bool = ...) -> None: ...
     class JSONElementType(TypeEngine):
         def string_bind_processor(self, dialect: Any): ...
@@ -293,10 +298,6 @@ class JSON(Indexable, TypeEngine[Union[str, Mapping, List]]):
     comparator_factory: Any = ...
     @property
     def python_type(self): ...
-    @property
-    def should_evaluate_none(self): ...
-    @should_evaluate_none.setter
-    def should_evaluate_none(self, value: Any) -> None: ...
     def bind_processor(self, dialect: Any): ...
     def result_processor(self, dialect: Any, coltype: Any): ...
 
diff --git a/sqlalchemy-stubs/sql/type_api.pyi b/sqlalchemy-stubs/sql/type_api.pyi
@@ -1,8 +1,15 @@
+from _typeshed import SupportsLessThan
 from typing import Any
+from typing import Callable
 from typing import Generic
+from typing import Mapping
 from typing import Optional
+from typing import Text
+from typing import Tuple
 from typing import Type
 from typing import TypeVar
+from typing import Union
+from typing_extensions import Protocol
 
 from . import operators as operators
 from .base import SchemaEventTarget as SchemaEventTarget
@@ -20,6 +27,14 @@ INDEXABLE: Any
 TABLEVALUE: Any
 
 _T = TypeVar("_T")
+_U = TypeVar("_U")
+
+_TE = TypeVar('_TE', bound=TypeEngine[Any])
+_NFE = TypeVar('_NFE', bound=NativeForEmulated)
+_TD = TypeVar('_TD', bound=TypeDecorator[Any])
+_VT = TypeVar('_VT', bound=Variant[Any])
+
+_SortKeyFunction = Callable[[Any], SupportsLessThan]
 
 class TypeEngine(Traversible, Generic[_T]):
     class Comparator(operators.ColumnOperators):
@@ -32,87 +47,85 @@ class TypeEngine(Traversible, Generic[_T]):
         def reverse_operate(self, op: Any, other: Any, **kwargs: Any): ...
         def __reduce__(self): ...
     hashable: bool = ...
-    comparator_factory: Any = ...
-    sort_key_function: Any = ...
+    comparator_factory: Type[Any] = ...
+    sort_key_function: Optional[_SortKeyFunction] = ...
     should_evaluate_none: bool = ...
-    def evaluates_none(self): ...
-    def copy(self, **kw: Any): ...
+    def evaluates_none(self: _TE) -> _TE: ...
+    def copy(self: _TE, **kw: Any) -> _TE: ...
     def compare_against_backend(
         self, dialect: Any, conn_type: Any
-    ) -> None: ...
-    def copy_value(self, value: Any): ...
-    def literal_processor(self, dialect: Any) -> None: ...
-    def bind_processor(self, dialect: Any) -> None: ...
-    def result_processor(self, dialect: Any, coltype: Any) -> None: ...
-    def column_expression(self, colexpr: Any) -> None: ...
-    def bind_expression(self, bindvalue: Any) -> None: ...
-    def compare_values(self, x: Any, y: Any): ...
-    def get_dbapi_type(self, dbapi: Any) -> None: ...
+    ) -> Any: ...
+    def copy_value(self, value: _T) -> _T: ...
+    def literal_processor(self, dialect: Any) -> Optional[Callable[..., Any]]: ...
+    def bind_processor(self, dialect: Any) -> Optional[Callable[..., Any]]: ...
+    def result_processor(self, dialect: Any, coltype: Any) -> Optional[Callable[..., Any]]: ...
+    def column_expression(self, colexpr: Any) -> Any: ...
+    def bind_expression(self, bindvalue: Any) -> Any: ...
+    def compare_values(self, x: Any, y: Any) -> bool: ...
+    def get_dbapi_type(self, dbapi: Any) -> Any: ...
     @property
     def python_type(self) -> Type[_T]: ...
-    def with_variant(self, type_: Any, dialect_name: Any): ...
-    def as_generic(self, allow_nulltype: bool = ...): ...
-    def dialect_impl(self, dialect: Any): ...
-    def adapt(self, cls: Any, **kw: Any): ...
-    def coerce_compared_value(self, op: Any, value: Any): ...
-    def compile(self, dialect: Optional[Any] = ...): ...
+    def with_variant(self, type_: Type[TypeEngine[_U]], dialect_name: str) -> Variant[_U]: ...
+    def as_generic(self, allow_nulltype: bool = ...) -> TypeEngine[Any]: ...
+    def dialect_impl(self, dialect: Any) -> Type[Any]: ...
+    def adapt(self, __cls: Type[_U], **kw: Any) -> _U: ...
+    def coerce_compared_value(self, op: Any, value: Any) -> TypeEngine[Any]: ...
+    def compile(self, dialect: Optional[Any] = ...) -> Any: ...
 
 class VisitableCheckKWArg(util.EnsureKWArgType, TraversibleType): ...
 
 class UserDefinedType:
     __visit_name__: str = ...
     ensure_kwarg: str = ...
-    def coerce_compared_value(self, op: Any, value: Any): ...
+    def coerce_compared_value(self, op: Any, value: Any) -> Any: ...
 
 class Emulated:
-    def adapt_to_emulated(self, impltype: Any, **kw: Any): ...
-    def adapt(self, impltype: Any, **kw: Any): ...
+    def adapt_to_emulated(self, impltype: Any, **kw: Any) -> Any: ...
+    def adapt(self, __impltype: Any, **kw: Any) -> Any: ...
 
 class NativeForEmulated:
     @classmethod
-    def adapt_native_to_emulated(cls, impl: Any, **kw: Any): ...
+    def adapt_native_to_emulated(cls, impl: Any, **kw: Any) -> Any: ...
     @classmethod
-    def adapt_emulated_to_native(cls, impl: Any, **kw: Any): ...
+    def adapt_emulated_to_native(cls: Type[_NFE], impl: Any, **kw: Any) -> _NFE: ...
 
-_TD = TypeVar("_TD")
-
-class TypeDecorator(SchemaEventTarget, TypeEngine[Any], Generic[_TD]):
+class TypeDecorator(SchemaEventTarget, TypeEngine[_T]):
     __visit_name__: str = ...
     impl: Any = ...
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
-    coerce_to_is_types: Any = ...
+    coerce_to_is_types: Tuple[Type[Any], ...] = ...
     class Comparator(TypeEngine.Comparator):
         def operate(self, op: Any, *other: Any, **kwargs: Any): ...
         def reverse_operate(self, op: Any, other: Any, **kwargs: Any): ...
     @property
-    def comparator_factory(self): ...
-    def type_engine(self, dialect: Any): ...
-    def load_dialect_impl(self, dialect: Any): ...
-    def __getattr__(self, key: Any): ...
-    def process_literal_param(self, value: Any, dialect: Any) -> None: ...
-    def process_bind_param(self, value: Any, dialect: Any) -> None: ...
-    def process_result_value(self, value: Any, dialect: Any) -> None: ...
-    def literal_processor(self, dialect: Any): ...
-    def bind_processor(self, dialect: Any): ...
-    def result_processor(self, dialect: Any, coltype: Any): ...
-    def bind_expression(self, bindparam: Any): ...
-    def column_expression(self, column: Any): ...
-    def coerce_compared_value(self, op: Any, value: Any): ...
-    def copy(self, **kw: Any): ...
-    def get_dbapi_type(self, dbapi: Any): ...
-    def compare_values(self, x: Any, y: Any): ...
+    def comparator_factory(self) -> Type[Any]: ...  # type: ignore[override]
+    def type_engine(self, dialect: Any) -> TypeEngine[Any]: ...
+    def load_dialect_impl(self, dialect: Any) -> TypeEngine[Any]: ...
+    def __getattr__(self, key: Any) -> Any: ...
+    def process_literal_param(self, value: Any, dialect: Any) -> Optional[str]: ...
+    def process_bind_param(self, value: Any, dialect: Any) -> Optional[Text]: ...
+    def process_result_value(self, value: Any, dialect: Any) -> Optional[_T]: ...
+    def literal_processor(self, dialect: Any) -> Callable[[Optional[_T]], Optional[str]]: ...
+    def bind_processor(self, dialect: Any) -> Callable[[Optional[_T]], Optional[str]]: ...
+    def result_processor(self, dialect: Any, coltype: Any) -> Callable[[Optional[Any]], Optional[_T]]: ...
+    def bind_expression(self, bindparam: Any) -> Any: ...
+    def column_expression(self, column: Any) -> Any: ...
+    def coerce_compared_value(self, op: Any, value: Any) -> Any: ...
+    def copy(self: _TD, **kw: Any) -> _TD: ...
+    def get_dbapi_type(self, dbapi: Any) -> Any: ...
+    def compare_values(self, x: Any, y: Any) -> bool: ...
     @property
-    def sort_key_function(self): ...
+    def sort_key_function(self) -> Optional[_SortKeyFunction]: ...  # type: ignore[override]
 
-class Variant(TypeDecorator[Any]):
-    impl: Any = ...
-    mapping: Any = ...
-    def __init__(self, base: Any, mapping: Any) -> None: ...
-    def coerce_compared_value(self, operator: Any, value: Any): ...
-    def load_dialect_impl(self, dialect: Any): ...
-    def with_variant(self, type_: Any, dialect_name: Any): ...
+class Variant(TypeDecorator[_T]):
+    impl: Type[TypeEngine[Any]] = ...
+    mapping: Mapping[str, TypeEngine[Any]] = ...
+    def __init__(self, base: Any, mapping: Mapping[str, TypeEngine[Any]]) -> None: ...
+    def coerce_compared_value(self: _VT, operator: Any, value: Any) -> Union[_VT, TypeEngine[Any]]: ...
+    def load_dialect_impl(self, dialect: Any) -> TypeEngine[Any]: ...
+    def with_variant(self, type_: Type[TypeEngine[_U]], dialect_name: str) -> Variant[_U]: ...
     @property
-    def comparator_factory(self): ...
+    def comparator_factory(self) -> Type[Any]: ...  # type: ignore[override]
 
 def to_instance(typeobj: Any, *arg: Any, **kw: Any): ...
 def adapt_type(typeobj: Any, colspecs: Any): ...
