establish mypy / typing approach for v2.0

large patch to get ORM / typing efforts started.
this is to support adding new test cases to mypy,
support dropping sqlalchemy2-stubs entirely from the
test suite, validate major ORM typing reorganization
to eliminate the need for the mypy plugin.

* New declarative approach which uses annotation
  introspection, fixes: #7535
* Mapped[] is now at the base of all ORM constructs
  that find themselves in classes, to support direct
  typing without plugins
* Mypy plugin updated for new typing structures
* Mypy test suite broken out into "plugin" tests vs.
  "plain" tests, and enhanced to better support test
  structures where we assert that various objects are
  introspected by the type checker as we expect.
  as we go forward with typing, we will
  add new use cases to "plain" where we can assert that
  types are introspected as we expect.
* For typing support, users will be much more exposed to the
  class names of things.  Add these all to "sqlalchemy" import
  space.
* Column(ForeignKey()) no longer needs to be `@declared_attr`
  if the FK refers to a remote table
* composite() attributes mapped to a dataclass no longer
  need to implement a `__composite_values__()` method
* with_variant() accepts multiple dialect names

Change-Id: I22797c0be73a8fbbd2d6f5e0c0b7258b17fe145d
Fixes: #7535
Fixes: #7551
References: #6810

9 files changed, 177 insertions, 54 deletions

diff --git a/lib/sqlalchemy/ext/associationproxy.py b/lib/sqlalchemy/ext/associationproxy.py
index e6a826c64..d5119907e 100644
--- a/lib/sqlalchemy/ext/associationproxy.py
+++ b/lib/sqlalchemy/ext/associationproxy.py
@@ -361,7 +361,7 @@ class AssociationProxyInstance:
         prop = orm.class_mapper(owning_class).get_property(target_collection)
 
         # this was never asserted before but this should be made clear.
-        if not isinstance(prop, orm.RelationshipProperty):
+        if not isinstance(prop, orm.Relationship):
             raise NotImplementedError(
                 "association proxy to a non-relationship "
                 "intermediary is not supported"
@@ -717,8 +717,8 @@ class AssociationProxyInstance:
         """Produce a proxied 'any' expression using EXISTS.
 
         This expression will be a composed product
-        using the :meth:`.RelationshipProperty.Comparator.any`
-        and/or :meth:`.RelationshipProperty.Comparator.has`
+        using the :meth:`.Relationship.Comparator.any`
+        and/or :meth:`.Relationship.Comparator.has`
         operators of the underlying proxied attributes.
 
         """
@@ -737,8 +737,8 @@ class AssociationProxyInstance:
         """Produce a proxied 'has' expression using EXISTS.
 
         This expression will be a composed product
-        using the :meth:`.RelationshipProperty.Comparator.any`
-        and/or :meth:`.RelationshipProperty.Comparator.has`
+        using the :meth:`.Relationship.Comparator.any`
+        and/or :meth:`.Relationship.Comparator.has`
         operators of the underlying proxied attributes.
 
         """
@@ -859,9 +859,9 @@ class ObjectAssociationProxyInstance(AssociationProxyInstance):
         """Produce a proxied 'contains' expression using EXISTS.
 
         This expression will be a composed product
-        using the :meth:`.RelationshipProperty.Comparator.any`,
-        :meth:`.RelationshipProperty.Comparator.has`,
-        and/or :meth:`.RelationshipProperty.Comparator.contains`
+        using the :meth:`.Relationship.Comparator.any`,
+        :meth:`.Relationship.Comparator.has`,
+        and/or :meth:`.Relationship.Comparator.contains`
         operators of the underlying proxied attributes.
         """
 
diff --git a/lib/sqlalchemy/ext/declarative/extensions.py b/lib/sqlalchemy/ext/declarative/extensions.py
index 5aff4dfe2..470ff6ad8 100644
--- a/lib/sqlalchemy/ext/declarative/extensions.py
+++ b/lib/sqlalchemy/ext/declarative/extensions.py
@@ -378,7 +378,7 @@ class DeferredReflection:
                 metadata = mapper.class_.metadata
                 for rel in mapper._props.values():
                     if (
-                        isinstance(rel, relationships.RelationshipProperty)
+                        isinstance(rel, relationships.Relationship)
                         and rel.secondary is not None
                     ):
                         if isinstance(rel.secondary, Table):
diff --git a/lib/sqlalchemy/ext/mypy/apply.py b/lib/sqlalchemy/ext/mypy/apply.py
index 99be194cd..4e244b5b9 100644
--- a/lib/sqlalchemy/ext/mypy/apply.py
+++ b/lib/sqlalchemy/ext/mypy/apply.py
@@ -36,6 +36,7 @@ from mypy.types import UnionType
 
 from . import infer
 from . import util
+from .names import expr_to_mapped_constructor
 from .names import NAMED_TYPE_SQLA_MAPPED
 
 
@@ -117,6 +118,7 @@ def re_apply_declarative_assignments(
         ):
 
             left_node = stmt.lvalues[0].node
+
             python_type_for_type = mapped_attr_lookup[
                 stmt.lvalues[0].name
             ].type
@@ -142,7 +144,7 @@ def re_apply_declarative_assignments(
                 )
             ):
 
-                python_type_for_type = (
+                new_python_type_for_type = (
                     infer.infer_type_from_right_hand_nameexpr(
                         api,
                         stmt,
@@ -152,19 +154,27 @@ def re_apply_declarative_assignments(
                     )
                 )
 
-                if python_type_for_type is None or isinstance(
-                    python_type_for_type, UnboundType
+                if new_python_type_for_type is not None and not isinstance(
+                    new_python_type_for_type, UnboundType
                 ):
-                    continue
+                    python_type_for_type = new_python_type_for_type
 
-                # update the SQLAlchemyAttribute with the better information
-                mapped_attr_lookup[
-                    stmt.lvalues[0].name
-                ].type = python_type_for_type
+                    # update the SQLAlchemyAttribute with the better
+                    # information
+                    mapped_attr_lookup[
+                        stmt.lvalues[0].name
+                    ].type = python_type_for_type
 
-                update_cls_metadata = True
+                    update_cls_metadata = True
 
-            if python_type_for_type is not None:
+            # for some reason if you have a Mapped type explicitly annotated,
+            # and here you set it again, mypy forgets how to do descriptors.
+            # no idea.  100% feeling around in the dark to see what sticks
+            if (
+                not isinstance(left_node.type, Instance)
+                or left_node.type.type.fullname != NAMED_TYPE_SQLA_MAPPED
+            ):
+                assert python_type_for_type is not None
                 left_node.type = api.named_type(
                     NAMED_TYPE_SQLA_MAPPED, [python_type_for_type]
                 )
@@ -202,6 +212,7 @@ def apply_type_to_mapped_statement(
     assert isinstance(left_node, Var)
 
     if left_hand_explicit_type is not None:
+        lvalue.is_inferred_def = False
         left_node.type = api.named_type(
             NAMED_TYPE_SQLA_MAPPED, [left_hand_explicit_type]
         )
@@ -224,7 +235,7 @@ def apply_type_to_mapped_statement(
     # _sa_Mapped._empty_constructor(<original CallExpr from rvalue>)
     # the original right-hand side is maintained so it gets type checked
     # internally
-    stmt.rvalue = util.expr_to_mapped_constructor(stmt.rvalue)
+    stmt.rvalue = expr_to_mapped_constructor(stmt.rvalue)
 
 
 def add_additional_orm_attributes(
diff --git a/lib/sqlalchemy/ext/mypy/decl_class.py b/lib/sqlalchemy/ext/mypy/decl_class.py
index c33c30e25..bd6c6f41e 100644
--- a/lib/sqlalchemy/ext/mypy/decl_class.py
+++ b/lib/sqlalchemy/ext/mypy/decl_class.py
@@ -337,7 +337,7 @@ def _scan_declarative_decorator_stmt(
     # <attr> : Mapped[<typ>] =
     # _sa_Mapped._empty_constructor(lambda: <function body>)
     # the function body is maintained so it gets type checked internally
-    rvalue = util.expr_to_mapped_constructor(
+    rvalue = names.expr_to_mapped_constructor(
         LambdaExpr(stmt.func.arguments, stmt.func.body)
     )
 
diff --git a/lib/sqlalchemy/ext/mypy/infer.py b/lib/sqlalchemy/ext/mypy/infer.py
index 3cd946e04..6a5e99e48 100644
--- a/lib/sqlalchemy/ext/mypy/infer.py
+++ b/lib/sqlalchemy/ext/mypy/infer.py
@@ -42,11 +42,13 @@ def infer_type_from_right_hand_nameexpr(
     left_hand_explicit_type: Optional[ProperType],
     infer_from_right_side: RefExpr,
 ) -> Optional[ProperType]:
-
     type_id = names.type_id_for_callee(infer_from_right_side)
-
     if type_id is None:
         return None
+    elif type_id is names.MAPPED:
+        python_type_for_type = _infer_type_from_mapped(
+            api, stmt, node, left_hand_explicit_type, infer_from_right_side
+        )
     elif type_id is names.COLUMN:
         python_type_for_type = _infer_type_from_decl_column(
             api, stmt, node, left_hand_explicit_type
@@ -245,7 +247,7 @@ def _infer_type_from_decl_composite_property(
     node: Var,
     left_hand_explicit_type: Optional[ProperType],
 ) -> Optional[ProperType]:
-    """Infer the type of mapping from a CompositeProperty."""
+    """Infer the type of mapping from a Composite."""
 
     assert isinstance(stmt.rvalue, CallExpr)
     target_cls_arg = stmt.rvalue.args[0]
@@ -271,6 +273,38 @@ def _infer_type_from_decl_composite_property(
         return python_type_for_type
 
 
+def _infer_type_from_mapped(
+    api: SemanticAnalyzerPluginInterface,
+    stmt: AssignmentStmt,
+    node: Var,
+    left_hand_explicit_type: Optional[ProperType],
+    infer_from_right_side: RefExpr,
+) -> Optional[ProperType]:
+    """Infer the type of mapping from a right side expression
+    that returns Mapped.
+
+
+    """
+    assert isinstance(stmt.rvalue, CallExpr)
+
+    # (Pdb) print(stmt.rvalue.callee)
+    # NameExpr(query_expression [sqlalchemy.orm._orm_constructors.query_expression])  # noqa: E501
+    # (Pdb) stmt.rvalue.callee.node
+    # <mypy.nodes.FuncDef object at 0x7f8d92fb5940>
+    # (Pdb) stmt.rvalue.callee.node.type
+    # def [_T] (default_expr: sqlalchemy.sql.elements.ColumnElement[_T`-1] =) -> sqlalchemy.orm.base.Mapped[_T`-1]  # noqa: E501
+    # sqlalchemy.orm.base.Mapped[_T`-1]
+    # the_mapped_type = stmt.rvalue.callee.node.type.ret_type
+
+    # TODO: look at generic ref and either use that,
+    # or reconcile w/ what's present, etc.
+    the_mapped_type = util.type_for_callee(infer_from_right_side)  # noqa
+
+    return infer_type_from_left_hand_type_only(
+        api, node, left_hand_explicit_type
+    )
+
+
 def _infer_type_from_decl_column_property(
     api: SemanticAnalyzerPluginInterface,
     stmt: AssignmentStmt,
diff --git a/lib/sqlalchemy/ext/mypy/names.py b/lib/sqlalchemy/ext/mypy/names.py
index b6f911979..ad4449e5b 100644
--- a/lib/sqlalchemy/ext/mypy/names.py
+++ b/lib/sqlalchemy/ext/mypy/names.py
@@ -12,11 +12,14 @@ from typing import Set
 from typing import Tuple
 from typing import Union
 
+from mypy.nodes import ARG_POS
+from mypy.nodes import CallExpr
 from mypy.nodes import ClassDef
 from mypy.nodes import Expression
 from mypy.nodes import FuncDef
 from mypy.nodes import MemberExpr
 from mypy.nodes import NameExpr
+from mypy.nodes import OverloadedFuncDef
 from mypy.nodes import SymbolNode
 from mypy.nodes import TypeAlias
 from mypy.nodes import TypeInfo
@@ -51,7 +54,7 @@ QUERY_EXPRESSION: int = util.symbol("QUERY_EXPRESSION")  # type: ignore
 NAMED_TYPE_BUILTINS_OBJECT = "builtins.object"
 NAMED_TYPE_BUILTINS_STR = "builtins.str"
 NAMED_TYPE_BUILTINS_LIST = "builtins.list"
-NAMED_TYPE_SQLA_MAPPED = "sqlalchemy.orm.attributes.Mapped"
+NAMED_TYPE_SQLA_MAPPED = "sqlalchemy.orm.base.Mapped"
 
 _lookup: Dict[str, Tuple[int, Set[str]]] = {
     "Column": (
@@ -61,11 +64,11 @@ _lookup: Dict[str, Tuple[int, Set[str]]] = {
             "sqlalchemy.sql.Column",
         },
     ),
-    "RelationshipProperty": (
+    "Relationship": (
         RELATIONSHIP,
         {
-            "sqlalchemy.orm.relationships.RelationshipProperty",
-            "sqlalchemy.orm.RelationshipProperty",
+            "sqlalchemy.orm.relationships.Relationship",
+            "sqlalchemy.orm.Relationship",
         },
     ),
     "registry": (
@@ -82,18 +85,18 @@ _lookup: Dict[str, Tuple[int, Set[str]]] = {
             "sqlalchemy.orm.ColumnProperty",
         },
     ),
-    "SynonymProperty": (
+    "Synonym": (
         SYNONYM_PROPERTY,
         {
-            "sqlalchemy.orm.descriptor_props.SynonymProperty",
-            "sqlalchemy.orm.SynonymProperty",
+            "sqlalchemy.orm.descriptor_props.Synonym",
+            "sqlalchemy.orm.Synonym",
         },
     ),
-    "CompositeProperty": (
+    "Composite": (
         COMPOSITE_PROPERTY,
         {
-            "sqlalchemy.orm.descriptor_props.CompositeProperty",
-            "sqlalchemy.orm.CompositeProperty",
+            "sqlalchemy.orm.descriptor_props.Composite",
+            "sqlalchemy.orm.Composite",
         },
     ),
     "MapperProperty": (
@@ -159,7 +162,10 @@ _lookup: Dict[str, Tuple[int, Set[str]]] = {
     ),
     "query_expression": (
         QUERY_EXPRESSION,
-        {"sqlalchemy.orm.query_expression"},
+        {
+            "sqlalchemy.orm.query_expression",
+            "sqlalchemy.orm._orm_constructors.query_expression",
+        },
     ),
 }
 
@@ -209,7 +215,19 @@ def type_id_for_unbound_type(
 
 def type_id_for_callee(callee: Expression) -> Optional[int]:
     if isinstance(callee, (MemberExpr, NameExpr)):
-        if isinstance(callee.node, FuncDef):
+        if isinstance(callee.node, OverloadedFuncDef):
+            if (
+                callee.node.impl
+                and callee.node.impl.type
+                and isinstance(callee.node.impl.type, CallableType)
+            ):
+                ret_type = get_proper_type(callee.node.impl.type.ret_type)
+
+                if isinstance(ret_type, Instance):
+                    return type_id_for_fullname(ret_type.type.fullname)
+
+            return None
+        elif isinstance(callee.node, FuncDef):
             if callee.node.type and isinstance(callee.node.type, CallableType):
                 ret_type = get_proper_type(callee.node.type.ret_type)
 
@@ -251,3 +269,15 @@ def type_id_for_fullname(fullname: str) -> Optional[int]:
         return type_id
     else:
         return None
+
+
+def expr_to_mapped_constructor(expr: Expression) -> CallExpr:
+    column_descriptor = NameExpr("__sa_Mapped")
+    column_descriptor.fullname = NAMED_TYPE_SQLA_MAPPED
+    member_expr = MemberExpr(column_descriptor, "_empty_constructor")
+    return CallExpr(
+        member_expr,
+        [expr],
+        [ARG_POS],
+        ["arg1"],
+    )
diff --git a/lib/sqlalchemy/ext/mypy/plugin.py b/lib/sqlalchemy/ext/mypy/plugin.py
index 0a21feb51..c9520fef3 100644
--- a/lib/sqlalchemy/ext/mypy/plugin.py
+++ b/lib/sqlalchemy/ext/mypy/plugin.py
@@ -40,6 +40,19 @@ from . import decl_class
 from . import names
 from . import util
 
+try:
+    import sqlalchemy_stubs  # noqa
+except ImportError:
+    pass
+else:
+    import sqlalchemy
+
+    raise ImportError(
+        f"The SQLAlchemy mypy plugin in SQLAlchemy "
+        f"{sqlalchemy.__version__} does not work with sqlalchemy-stubs or "
+        "sqlalchemy2-stubs installed"
+    )
+
 
 class SQLAlchemyPlugin(Plugin):
     def get_dynamic_class_hook(
diff --git a/lib/sqlalchemy/ext/mypy/util.py b/lib/sqlalchemy/ext/mypy/util.py
index fa42074c3..741772eac 100644
--- a/lib/sqlalchemy/ext/mypy/util.py
+++ b/lib/sqlalchemy/ext/mypy/util.py
@@ -10,24 +10,27 @@ from typing import Type as TypingType
 from typing import TypeVar
 from typing import Union
 
-from mypy.nodes import ARG_POS
 from mypy.nodes import CallExpr
 from mypy.nodes import ClassDef
 from mypy.nodes import CLASSDEF_NO_INFO
 from mypy.nodes import Context
 from mypy.nodes import Expression
+from mypy.nodes import FuncDef
 from mypy.nodes import IfStmt
 from mypy.nodes import JsonDict
 from mypy.nodes import MemberExpr
 from mypy.nodes import NameExpr
 from mypy.nodes import Statement
 from mypy.nodes import SymbolTableNode
+from mypy.nodes import TypeAlias
 from mypy.nodes import TypeInfo
 from mypy.plugin import ClassDefContext
 from mypy.plugin import DynamicClassDefContext
 from mypy.plugin import SemanticAnalyzerPluginInterface
 from mypy.plugins.common import deserialize_and_fixup_type
 from mypy.typeops import map_type_from_supertype
+from mypy.types import CallableType
+from mypy.types import get_proper_type
 from mypy.types import Instance
 from mypy.types import NoneType
 from mypy.types import Type
@@ -231,6 +234,25 @@ def flatten_typechecking(stmts: Iterable[Statement]) -> Iterator[Statement]:
             yield stmt
 
 
+def type_for_callee(callee: Expression) -> Optional[Union[Instance, TypeInfo]]:
+    if isinstance(callee, (MemberExpr, NameExpr)):
+        if isinstance(callee.node, FuncDef):
+            if callee.node.type and isinstance(callee.node.type, CallableType):
+                ret_type = get_proper_type(callee.node.type.ret_type)
+
+                if isinstance(ret_type, Instance):
+                    return ret_type
+
+            return None
+        elif isinstance(callee.node, TypeAlias):
+            target_type = get_proper_type(callee.node.target)
+            if isinstance(target_type, Instance):
+                return target_type
+        elif isinstance(callee.node, TypeInfo):
+            return callee.node
+    return None
+
+
 def unbound_to_instance(
     api: SemanticAnalyzerPluginInterface, typ: Type
 ) -> Type:
@@ -290,15 +312,3 @@ def info_for_cls(
         return sym.node
 
     return cls.info
-
-
-def expr_to_mapped_constructor(expr: Expression) -> CallExpr:
-    column_descriptor = NameExpr("__sa_Mapped")
-    column_descriptor.fullname = "sqlalchemy.orm.attributes.Mapped"
-    member_expr = MemberExpr(column_descriptor, "_empty_constructor")
-    return CallExpr(
-        member_expr,
-        [expr],
-        [ARG_POS],
-        ["arg1"],
-    )
diff --git a/lib/sqlalchemy/ext/orderinglist.py b/lib/sqlalchemy/ext/orderinglist.py
index 5a327d1a5..5384851b1 100644
--- a/lib/sqlalchemy/ext/orderinglist.py
+++ b/lib/sqlalchemy/ext/orderinglist.py
@@ -119,14 +119,28 @@ start numbering at 1 or some other integer, provide ``count_from=1``.
 
 
 """
+from typing import Callable
+from typing import List
+from typing import Optional
+from typing import Sequence
+from typing import TypeVar
+
 from ..orm.collections import collection
 from ..orm.collections import collection_adapter
 
+_T = TypeVar("_T")
+OrderingFunc = Callable[[int, Sequence[_T]], int]
+
 
 __all__ = ["ordering_list"]
 
 
-def ordering_list(attr, count_from=None, **kw):
+def ordering_list(
+    attr: str,
+    count_from: Optional[int] = None,
+    ordering_func: Optional[OrderingFunc] = None,
+    reorder_on_append: bool = False,
+) -> Callable[[], "OrderingList"]:
     """Prepares an :class:`OrderingList` factory for use in mapper definitions.
 
     Returns an object suitable for use as an argument to a Mapper
@@ -157,7 +171,11 @@ def ordering_list(attr, count_from=None, **kw):
 
     """
 
-    kw = _unsugar_count_from(count_from=count_from, **kw)
+    kw = _unsugar_count_from(
+        count_from=count_from,
+        ordering_func=ordering_func,
+        reorder_on_append=reorder_on_append,
+    )
     return lambda: OrderingList(attr, **kw)
 
 
@@ -207,7 +225,7 @@ def _unsugar_count_from(**kw):
     return kw
 
 
-class OrderingList(list):
+class OrderingList(List[_T]):
     """A custom list that manages position information for its children.
 
     The :class:`.OrderingList` object is normally set up using the
@@ -216,8 +234,15 @@ class OrderingList(list):
 
     """
 
+    ordering_attr: str
+    ordering_func: OrderingFunc
+    reorder_on_append: bool
+
     def __init__(
-        self, ordering_attr=None, ordering_func=None, reorder_on_append=False
+        self,
+        ordering_attr: Optional[str] = None,
+        ordering_func: Optional[OrderingFunc] = None,
+        reorder_on_append: bool = False,
     ):
         """A custom list that manages position information for its children.
 
@@ -282,7 +307,7 @@ class OrderingList(list):
     def _set_order_value(self, entity, value):
         setattr(entity, self.ordering_attr, value)
 
-    def reorder(self):
+    def reorder(self) -> None:
         """Synchronize ordering for the entire collection.
 
         Sweeps through the list and ensures that each object has accurate