Merge pull request #6 from bryanforbes/improve-sql-type-api

zzzeek · web-flow · commit 14ecd6de9738 · 2021-03-25T10:19:37.000-04:00
Improve typing for sql.type_api classes
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
 
+_T = TypeVar("_T")
+
 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[_T], **kw: Any) -> _T: ...
     @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[_T], **kw: Any) -> _T: ...
     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): ...
 
@@ -327,7 +328,6 @@ class ARRAY(SchemaEventTarget, Indexable, Concatenable, TypeEngine[List]):
 class TupleType(TypeEngine[TupleType]):
     types: Any = ...
     def __init__(self, *types: Any) -> None: ...
-    def result_processor(self, dialect: Any, coltype: Any) -> None: ...
 
 class REAL(Float):
     __visit_name__: str = ...
diff --git a/sqlalchemy-stubs/sql/type_api.pyi b/sqlalchemy-stubs/sql/type_api.pyi
@@ -1,8 +1,15 @@
+import sys
 from typing import Any
+from typing import Callable
 from typing import Generic
+from typing import Mapping
 from typing import Optional
+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,30 @@ INDEXABLE: Any
 TABLEVALUE: Any
 
 _T = TypeVar("_T")
+_T_co = TypeVar("_T_co", covariant=True)
+_T_contra = TypeVar("_T_contra", contravariant=True)
+_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])
+
+if sys.version_info[0] < 3:
+    from _typeshed import SupportsLessThan
+
+    _SortKeyFunction = Callable[[Any], SupportsLessThan]
+else:
+    _SortKeyFunction = Callable[[Any], Any]
+
+class _LiteralProcessor(Protocol[_T_contra]):
+    def __call__(self, __value: Optional[_T_contra]) -> str: ...
+
+class _BindProcessor(Protocol[_T_contra]):
+    def __call__(self, __value: Optional[_T_contra]) -> Optional[Any]: ...
+
+class _ResultProcessor(Protocol[_T_co]):
+    def __call__(self, __value: Optional[Any]) -> Optional[_T_co]: ...
 
 class TypeEngine(Traversible, Generic[_T]):
     class Comparator(operators.ColumnOperators):
@@ -32,87 +63,105 @@ 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 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: ...
+    def evaluates_none(self: _TE) -> _TE: ...
+    def copy(self: _TE, **kw: Any) -> _TE: ...
+    def compare_against_backend(self, dialect: Any, conn_type: Any) -> Any: ...
+    def copy_value(self, value: _T) -> _T: ...
+    def literal_processor(
+        self, dialect: Any
+    ) -> Optional[_LiteralProcessor[_T]]: ...
+    def bind_processor(self, dialect: Any) -> Optional[_BindProcessor[_T]]: ...
+    def result_processor(
+        self, dialect: Any, coltype: Any
+    ) -> Optional[_ResultProcessor[_T]]: ...
+    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): ...
-
-_TD = TypeVar("_TD")
+    def adapt_emulated_to_native(
+        cls: Type[_NFE], impl: Any, **kw: Any
+    ) -> _NFE: ...
 
-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: Optional[_T], dialect: Any
+    ) -> str: ...
+    def process_bind_param(self, value: Optional[_T], dialect: Any) -> Any: ...
+    def process_result_value(
+        self, value: Any, dialect: Any
+    ) -> Optional[_T]: ...
+    def literal_processor(self, dialect: Any) -> _LiteralProcessor[_T]: ...
+    def bind_processor(self, dialect: Any) -> _BindProcessor[_T]: ...
+    def result_processor(
+        self, dialect: Any, coltype: Any
+    ) -> _ResultProcessor[_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: 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): ...
