Make TypeEngine and RelationshipProperty covariant (#45)

Author: ilevkivskyi

sqlalchemy-stubs/orm/relationships.pyi

@@ -12,10 +12,14 @@ def foreign(expr): ...
 
 
 _T = TypeVar('_T')
+_T_co = TypeVar('_T_co', covariant=True)
+
 
 # Note: typical use case is where argument is a string, so this will require
 # a plugin to infer '_T', otherwise a user will need to write an explicit annotation.
-class RelationshipProperty(StrategizedProperty, Generic[_T]):
+# It is not clear whether RelationshipProperty is covariant at this stage since
+# many types are still missing.
+class RelationshipProperty(StrategizedProperty, Generic[_T_co]):
     strategy_wildcard_key: str = ...
     uselist: Any = ...
     argument: Any = ...
@@ -93,9 +97,9 @@ class RelationshipProperty(StrategizedProperty, Generic[_T]):
     def do_init(self): ...
     # This doesn't exist at runtime, but is present here for better typing.
     @overload
-    def __get__(self, instance: None, owner: Any) -> RelationshipProperty[_T]: ...
+    def __get__(self, instance: None, owner: Any) -> RelationshipProperty[_T_co]: ...
     @overload
-    def __get__(self, instance: object, owner: Any) -> _T: ...
+    def __get__(self, instance: object, owner: Any) -> _T_co: ...
 
 class JoinCondition(object):
     parent_selectable: Any = ...

sqlalchemy-stubs/sql/type_api.pyi

@@ -6,6 +6,7 @@ from . import operators
 from .sqltypes import NullType, Integer, Boolean, String, MatchType, Indexable
 
 _T = TypeVar('_T')
+_T_co = TypeVar('_T_co', covariant=True)
 
 # TODO: simplify import cycle with sqltypes
 BOOLEANTYPE: Boolean = ...
@@ -15,7 +16,9 @@ STRINGTYPE: String = ...
 MATCHTYPE: MatchType = ...
 INDEXABLE: Indexable = ...
 
-class TypeEngine(Visitable, Generic[_T]):
+# It is not 100% clear that TypeEngine is covariant at this stage since
+# many types are still missing. Anyway, this looks plausible and practical.
+class TypeEngine(Visitable, Generic[_T_co]):
     class Comparator(operators.ColumnOperators):
         default_comparator: Any = ...
         expr: Any = ...
@@ -39,7 +42,7 @@ class TypeEngine(Visitable, Generic[_T]):
     def compare_values(self, x, y): ...
     def get_dbapi_type(self, dbapi): ...
     @property
-    def python_type(self) -> Type[_T]: ...
+    def python_type(self) -> Type[_T_co]: ...
     def with_variant(self, type_, dialect_name): ...
     def dialect_impl(self, dialect): ...
     def adapt(self, *args, **kw): ...

test/test-data/sqlalchemy-basics.test

@@ -16,6 +16,18 @@ reveal_type(user.id)  # E: Revealed type is 'builtins.int*'
 reveal_type(User.name)  # E: Revealed type is 'sqlalchemy.sql.schema.Column[builtins.str*]'
 [out]
 
+[case testTypeEngineCovariance]
+from sqlalchemy import Column, Integer, String
+from sqlalchemy.sql.type_api import TypeEngine
+
+from typing import TypeVar, Optional
+T = TypeVar('T', bound=int)
+
+def func(tp: TypeEngine[Optional[T]]) -> T: ...
+reveal_type(func(Integer()))  # E: Revealed type is 'builtins.int*'
+func(String())  # E: Value of type variable "T" of "func" cannot be "str"
+[out]
+
 [case testColumnFieldsInferredInstance]
 from typing import Any