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

1 files changed, 232 insertions, 0 deletions

diff --git a/lib/sqlalchemy/orm/mapped_collection.py b/lib/sqlalchemy/orm/mapped_collection.py
new file mode 100644
index 000000000..75abeef4c
--- /dev/null
+++ b/lib/sqlalchemy/orm/mapped_collection.py
@@ -0,0 +1,232 @@
+# orm/collections.py
+# Copyright (C) 2005-2022 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+
+import operator
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Type
+from typing import TypeVar
+
+from . import base
+from .collections import collection
+from .. import exc as sa_exc
+from .. import util
+from ..sql import coercions
+from ..sql import expression
+from ..sql import roles
+
+_KT = TypeVar("_KT", bound=Any)
+_VT = TypeVar("_VT", bound=Any)
+
+
+class _PlainColumnGetter:
+    """Plain column getter, stores collection of Column objects
+    directly.
+
+    Serializes to a :class:`._SerializableColumnGetterV2`
+    which has more expensive __call__() performance
+    and some rare caveats.
+
+    """
+
+    __slots__ = ("cols", "composite")
+
+    def __init__(self, cols):
+        self.cols = cols
+        self.composite = len(cols) > 1
+
+    def __reduce__(self):
+        return _SerializableColumnGetterV2._reduce_from_cols(self.cols)
+
+    def _cols(self, mapper):
+        return self.cols
+
+    def __call__(self, value):
+        state = base.instance_state(value)
+        m = base._state_mapper(state)
+
+        key = [
+            m._get_state_attr_by_column(state, state.dict, col)
+            for col in self._cols(m)
+        ]
+
+        if self.composite:
+            return tuple(key)
+        else:
+            return key[0]
+
+
+class _SerializableColumnGetterV2(_PlainColumnGetter):
+    """Updated serializable getter which deals with
+    multi-table mapped classes.
+
+    Two extremely unusual cases are not supported.
+    Mappings which have tables across multiple metadata
+    objects, or which are mapped to non-Table selectables
+    linked across inheriting mappers may fail to function
+    here.
+
+    """
+
+    __slots__ = ("colkeys",)
+
+    def __init__(self, colkeys):
+        self.colkeys = colkeys
+        self.composite = len(colkeys) > 1
+
+    def __reduce__(self):
+        return self.__class__, (self.colkeys,)
+
+    @classmethod
+    def _reduce_from_cols(cls, cols):
+        def _table_key(c):
+            if not isinstance(c.table, expression.TableClause):
+                return None
+            else:
+                return c.table.key
+
+        colkeys = [(c.key, _table_key(c)) for c in cols]
+        return _SerializableColumnGetterV2, (colkeys,)
+
+    def _cols(self, mapper):
+        cols = []
+        metadata = getattr(mapper.local_table, "metadata", None)
+        for (ckey, tkey) in self.colkeys:
+            if tkey is None or metadata is None or tkey not in metadata:
+                cols.append(mapper.local_table.c[ckey])
+            else:
+                cols.append(metadata.tables[tkey].c[ckey])
+        return cols
+
+
+def column_mapped_collection(mapping_spec):
+    """A dictionary-based collection type with column-based keying.
+
+    Returns a :class:`.MappedCollection` factory with a keying function
+    generated from mapping_spec, which may be a Column or a sequence
+    of Columns.
+
+    The key value must be immutable for the lifetime of the object.  You
+    can not, for example, map on foreign key values if those key values will
+    change during the session, i.e. from None to a database-assigned integer
+    after a session flush.
+
+    """
+    cols = [
+        coercions.expect(roles.ColumnArgumentRole, q, argname="mapping_spec")
+        for q in util.to_list(mapping_spec)
+    ]
+    keyfunc = _PlainColumnGetter(cols)
+    return _mapped_collection_cls(keyfunc)
+
+
+def attribute_mapped_collection(attr_name: str) -> Type["MappedCollection"]:
+    """A dictionary-based collection type with attribute-based keying.
+
+    Returns a :class:`.MappedCollection` factory with a keying based on the
+    'attr_name' attribute of entities in the collection, where ``attr_name``
+    is the string name of the attribute.
+
+    .. warning:: the key value must be assigned to its final value
+       **before** it is accessed by the attribute mapped collection.
+       Additionally, changes to the key attribute are **not tracked**
+       automatically, which means the key in the dictionary is not
+       automatically synchronized with the key value on the target object
+       itself.  See the section :ref:`key_collections_mutations`
+       for an example.
+
+    """
+    getter = operator.attrgetter(attr_name)
+    return _mapped_collection_cls(getter)
+
+
+def mapped_collection(
+    keyfunc: Callable[[Any], _KT]
+) -> Type["MappedCollection[_KT, Any]"]:
+    """A dictionary-based collection type with arbitrary keying.
+
+    Returns a :class:`.MappedCollection` factory with a keying function
+    generated from keyfunc, a callable that takes an entity and returns a
+    key value.
+
+    The key value must be immutable for the lifetime of the object.  You
+    can not, for example, map on foreign key values if those key values will
+    change during the session, i.e. from None to a database-assigned integer
+    after a session flush.
+
+    """
+    return _mapped_collection_cls(keyfunc)
+
+
+class MappedCollection(Dict[_KT, _VT]):
+    """A basic dictionary-based collection class.
+
+    Extends dict with the minimal bag semantics that collection
+    classes require. ``set`` and ``remove`` are implemented in terms
+    of a keying function: any callable that takes an object and
+    returns an object for use as a dictionary key.
+
+    """
+
+    def __init__(self, keyfunc):
+        """Create a new collection with keying provided by keyfunc.
+
+        keyfunc may be any callable that takes an object and returns an object
+        for use as a dictionary key.
+
+        The keyfunc will be called every time the ORM needs to add a member by
+        value-only (such as when loading instances from the database) or
+        remove a member.  The usual cautions about dictionary keying apply-
+        ``keyfunc(object)`` should return the same output for the life of the
+        collection.  Keying based on mutable properties can result in
+        unreachable instances "lost" in the collection.
+
+        """
+        self.keyfunc = keyfunc
+
+    @classmethod
+    def _unreduce(cls, keyfunc, values):
+        mp = MappedCollection(keyfunc)
+        mp.update(values)
+        return mp
+
+    def __reduce__(self):
+        return (MappedCollection._unreduce, (self.keyfunc, dict(self)))
+
+    @collection.appender
+    @collection.internally_instrumented
+    def set(self, value, _sa_initiator=None):
+        """Add an item by value, consulting the keyfunc for the key."""
+
+        key = self.keyfunc(value)
+        self.__setitem__(key, value, _sa_initiator)
+
+    @collection.remover
+    @collection.internally_instrumented
+    def remove(self, value, _sa_initiator=None):
+        """Remove an item by value, consulting the keyfunc for the key."""
+
+        key = self.keyfunc(value)
+        # Let self[key] raise if key is not in this collection
+        # testlib.pragma exempt:__ne__
+        if self[key] != value:
+            raise sa_exc.InvalidRequestError(
+                "Can not remove '%s': collection holds '%s' for key '%s'. "
+                "Possible cause: is the MappedCollection key function "
+                "based on mutable properties or properties that only obtain "
+                "values after flush?" % (value, self[key], key)
+            )
+        self.__delitem__(key, _sa_initiator)
+
+
+def _mapped_collection_cls(keyfunc):
+    class _MKeyfuncMapped(MappedCollection):
+        def __init__(self):
+            super().__init__(keyfunc)
+
+    return _MKeyfuncMapped