Merge "revenge of pep 484" into main

1 files changed, 147 insertions, 82 deletions

diff --git a/lib/sqlalchemy/orm/decl_api.py b/lib/sqlalchemy/orm/decl_api.py
index fbe35f92a..1c343b04c 100644
--- a/lib/sqlalchemy/orm/decl_api.py
+++ b/lib/sqlalchemy/orm/decl_api.py
@@ -4,6 +4,7 @@
 #
 # This module is part of SQLAlchemy and is released under
 # the MIT License: https://www.opensource.org/licenses/mit-license.php
+
 """Public API functions and helpers for declarative."""
 
 from __future__ import annotations
@@ -14,16 +15,21 @@ import typing
 from typing import Any
 from typing import Callable
 from typing import ClassVar
+from typing import Dict
+from typing import FrozenSet
+from typing import Iterator
 from typing import Mapping
 from typing import Optional
+from typing import overload
+from typing import Set
 from typing import Type
+from typing import TYPE_CHECKING
 from typing import TypeVar
 from typing import Union
 import weakref
 
 from . import attributes
 from . import clsregistry
-from . import exc as orm_exc
 from . import instrumentation
 from . import interfaces
 from . import mapperlib
@@ -38,24 +44,40 @@ from .decl_base import _del_attribute
 from .decl_base import _mapper
 from .descriptor_props import Synonym as _orm_synonym
 from .mapper import Mapper
+from .state import InstanceState
 from .. import exc
 from .. import inspection
 from .. import util
 from ..sql.elements import SQLCoreOperations
 from ..sql.schema import MetaData
 from ..sql.selectable import FromClause
-from ..sql.type_api import TypeEngine
 from ..util import hybridmethod
 from ..util import hybridproperty
 from ..util import typing as compat_typing
+from ..util.typing import CallableReference
+from ..util.typing import Literal
 
+if TYPE_CHECKING:
+    from ._typing import _O
+    from ._typing import _RegistryType
+    from .descriptor_props import Synonym
+    from .instrumentation import ClassManager
+    from .interfaces import MapperProperty
+    from ..sql._typing import _TypeEngineArgument
 
 _T = TypeVar("_T", bound=Any)
 
-_TypeAnnotationMapType = Mapping[Type, Union[Type[TypeEngine], TypeEngine]]
+# it's not clear how to have Annotated, Union objects etc. as keys here
+# from a typing perspective so just leave it open ended for now
+_TypeAnnotationMapType = Mapping[Any, "_TypeEngineArgument[Any]"]
+_MutableTypeAnnotationMapType = Dict[Any, "_TypeEngineArgument[Any]"]
+
+_DeclaredAttrDecorated = Callable[
+    ..., Union[Mapped[_T], SQLCoreOperations[_T]]
+]
 
 
-def has_inherited_table(cls):
+def has_inherited_table(cls: Type[_O]) -> bool:
     """Given a class, return True if any of the classes it inherits from has a
     mapped table, otherwise return False.
 
@@ -75,13 +97,13 @@ def has_inherited_table(cls):
 
 
 class _DynamicAttributesType(type):
-    def __setattr__(cls, key, value):
+    def __setattr__(cls, key: str, value: Any) -> None:
         if "__mapper__" in cls.__dict__:
             _add_attribute(cls, key, value)
         else:
             type.__setattr__(cls, key, value)
 
-    def __delattr__(cls, key):
+    def __delattr__(cls, key: str) -> None:
         if "__mapper__" in cls.__dict__:
             _del_attribute(cls, key)
         else:
@@ -89,7 +111,7 @@ class _DynamicAttributesType(type):
 
 
 class DeclarativeAttributeIntercept(
-    _DynamicAttributesType, inspection.Inspectable["Mapper[Any]"]
+    _DynamicAttributesType, inspection.Inspectable[Mapper[Any]]
 ):
     """Metaclass that may be used in conjunction with the
     :class:`_orm.DeclarativeBase` class to support addition of class
@@ -99,10 +121,10 @@ class DeclarativeAttributeIntercept(
 
 
 class DeclarativeMeta(
-    _DynamicAttributesType, inspection.Inspectable["Mapper[Any]"]
+    _DynamicAttributesType, inspection.Inspectable[Mapper[Any]]
 ):
     metadata: MetaData
-    registry: "RegistryType"
+    registry: RegistryType
 
     def __init__(
         cls, classname: Any, bases: Any, dict_: Any, **kw: Any
@@ -130,7 +152,9 @@ class DeclarativeMeta(
         type.__init__(cls, classname, bases, dict_)
 
 
-def synonym_for(name, map_column=False):
+def synonym_for(
+    name: str, map_column: bool = False
+) -> Callable[[Callable[..., Any]], Synonym[Any]]:
     """Decorator that produces an :func:`_orm.synonym`
     attribute in conjunction with a Python descriptor.
 
@@ -164,7 +188,7 @@ def synonym_for(name, map_column=False):
 
     """
 
-    def decorate(fn):
+    def decorate(fn: Callable[..., Any]) -> Synonym[Any]:
         return _orm_synonym(name, map_column=map_column, descriptor=fn)
 
     return decorate
@@ -255,16 +279,16 @@ class declared_attr(interfaces._MappedAttribute[_T]):
 
     if typing.TYPE_CHECKING:
 
-        def __set__(self, instance, value):
+        def __set__(self, instance: Any, value: Any) -> None:
             ...
 
-        def __delete__(self, instance: Any):
+        def __delete__(self, instance: Any) -> None:
             ...
 
     def __init__(
         self,
-        fn: Callable[..., Union[Mapped[_T], SQLCoreOperations[_T]]],
-        cascading=False,
+        fn: _DeclaredAttrDecorated[_T],
+        cascading: bool = False,
     ):
         self.fget = fn
         self._cascading = cascading
@@ -273,10 +297,28 @@ class declared_attr(interfaces._MappedAttribute[_T]):
     def _collect_return_annotation(self) -> Optional[Type[Any]]:
         return util.get_annotations(self.fget).get("return")
 
-    def __get__(self, instance, owner) -> InstrumentedAttribute[_T]:
+    # this is the Mapped[] API where at class descriptor get time we want
+    # the type checker to see InstrumentedAttribute[_T].   However the
+    # callable function prior to mapping in fact calls the given
+    # declarative function that does not return InstrumentedAttribute
+    @overload
+    def __get__(self, instance: None, owner: Any) -> InstrumentedAttribute[_T]:
+        ...
+
+    @overload
+    def __get__(self, instance: object, owner: Any) -> _T:
+        ...
+
+    def __get__(
+        self, instance: Optional[object], owner: Any
+    ) -> Union[InstrumentedAttribute[_T], _T]:
         # the declared_attr needs to make use of a cache that exists
         # for the span of the declarative scan_attributes() phase.
         # to achieve this we look at the class manager that's configured.
+
+        # note this method should not be called outside of the declarative
+        # setup phase
+
         cls = owner
         manager = attributes.opt_manager_of_class(cls)
         if manager is None:
@@ -287,30 +329,33 @@ class declared_attr(interfaces._MappedAttribute[_T]):
                     "Unmanaged access of declarative attribute %s from "
                     "non-mapped class %s" % (self.fget.__name__, cls.__name__)
                 )
-            return self.fget(cls)
+            return self.fget(cls)  # type: ignore
         elif manager.is_mapped:
             # the class is mapped, which means we're outside of the declarative
             # scan setup, just run the function.
-            return self.fget(cls)
+            return self.fget(cls)  # type: ignore
 
         # here, we are inside of the declarative scan.  use the registry
         # that is tracking the values of these attributes.
         declarative_scan = manager.declarative_scan()
+
+        # assert that we are in fact in the declarative scan
         assert declarative_scan is not None
+
         reg = declarative_scan.declared_attr_reg
 
         if self in reg:
-            return reg[self]
+            return reg[self]  # type: ignore
         else:
             reg[self] = obj = self.fget(cls)
-            return obj
+            return obj  # type: ignore
 
     @hybridmethod
-    def _stateful(cls, **kw):
+    def _stateful(cls, **kw: Any) -> _stateful_declared_attr[_T]:
         return _stateful_declared_attr(**kw)
 
     @hybridproperty
-    def cascading(cls):
+    def cascading(cls) -> _stateful_declared_attr[_T]:
         """Mark a :class:`.declared_attr` as cascading.
 
         This is a special-use modifier which indicates that a column
@@ -372,20 +417,23 @@ class declared_attr(interfaces._MappedAttribute[_T]):
         return cls._stateful(cascading=True)
 
 
-class _stateful_declared_attr(declared_attr):
-    def __init__(self, **kw):
+class _stateful_declared_attr(declared_attr[_T]):
+    kw: Dict[str, Any]
+
+    def __init__(self, **kw: Any):
         self.kw = kw
 
-    def _stateful(self, **kw):
+    @hybridmethod
+    def _stateful(self, **kw: Any) -> _stateful_declared_attr[_T]:
         new_kw = self.kw.copy()
         new_kw.update(kw)
         return _stateful_declared_attr(**new_kw)
 
-    def __call__(self, fn):
+    def __call__(self, fn: _DeclaredAttrDecorated[_T]) -> declared_attr[_T]:
         return declared_attr(fn, **self.kw)
 
 
-def declarative_mixin(cls):
+def declarative_mixin(cls: Type[_T]) -> Type[_T]:
     """Mark a class as providing the feature of "declarative mixin".
 
     E.g.::
@@ -427,9 +475,9 @@ def declarative_mixin(cls):
     return cls
 
 
-def _setup_declarative_base(cls):
+def _setup_declarative_base(cls: Type[Any]) -> None:
     if "metadata" in cls.__dict__:
-        metadata = cls.metadata
+        metadata = cls.metadata  # type: ignore
     else:
         metadata = None
 
@@ -457,15 +505,15 @@ def _setup_declarative_base(cls):
         reg = registry(
             metadata=metadata, type_annotation_map=type_annotation_map
         )
-        cls.registry = reg
+        cls.registry = reg  # type: ignore
 
-    cls._sa_registry = reg
+    cls._sa_registry = reg  # type: ignore
 
     if "metadata" not in cls.__dict__:
-        cls.metadata = cls.registry.metadata
+        cls.metadata = cls.registry.metadata  # type: ignore
 
 
-class DeclarativeBaseNoMeta(inspection.Inspectable["Mapper"]):
+class DeclarativeBaseNoMeta(inspection.Inspectable[Mapper[Any]]):
     """Same as :class:`_orm.DeclarativeBase`, but does not use a metaclass
     to intercept new attributes.
 
@@ -477,10 +525,10 @@ class DeclarativeBaseNoMeta(inspection.Inspectable["Mapper"]):
 
     """
 
-    registry: ClassVar["registry"]
-    _sa_registry: ClassVar["registry"]
+    registry: ClassVar[_RegistryType]
+    _sa_registry: ClassVar[_RegistryType]
     metadata: ClassVar[MetaData]
-    __mapper__: ClassVar[Mapper]
+    __mapper__: ClassVar[Mapper[Any]]
     __table__: Optional[FromClause]
 
     if typing.TYPE_CHECKING:
@@ -496,7 +544,7 @@ class DeclarativeBaseNoMeta(inspection.Inspectable["Mapper"]):
 
 
 class DeclarativeBase(
-    inspection.Inspectable["InstanceState"],
+    inspection.Inspectable[InstanceState[Any]],
     metaclass=DeclarativeAttributeIntercept,
 ):
     """Base class used for declarative class definitions.
@@ -557,10 +605,10 @@ class DeclarativeBase(
 
     """
 
-    registry: ClassVar["registry"]
-    _sa_registry: ClassVar["registry"]
+    registry: ClassVar[_RegistryType]
+    _sa_registry: ClassVar[_RegistryType]
     metadata: ClassVar[MetaData]
-    __mapper__: ClassVar[Mapper]
+    __mapper__: ClassVar[Mapper[Any]]
     __table__: Optional[FromClause]
 
     if typing.TYPE_CHECKING:
@@ -572,10 +620,12 @@ class DeclarativeBase(
         if DeclarativeBase in cls.__bases__:
             _setup_declarative_base(cls)
         else:
-            cls._sa_registry.map_declaratively(cls)
+            _as_declarative(cls._sa_registry, cls, cls.__dict__)
 
 
-def add_mapped_attribute(target, key, attr):
+def add_mapped_attribute(
+    target: Type[_O], key: str, attr: MapperProperty[Any]
+) -> None:
     """Add a new mapped attribute to an ORM mapped class.
 
     E.g.::
@@ -593,14 +643,15 @@ def add_mapped_attribute(target, key, attr):
 
 
 def declarative_base(
+    *,
     metadata: Optional[MetaData] = None,
-    mapper=None,
-    cls=object,
-    name="Base",
+    mapper: Optional[Callable[..., Mapper[Any]]] = None,
+    cls: Type[Any] = object,
+    name: str = "Base",
     class_registry: Optional[clsregistry._ClsRegistryType] = None,
     type_annotation_map: Optional[_TypeAnnotationMapType] = None,
     constructor: Callable[..., None] = _declarative_constructor,
-    metaclass=DeclarativeMeta,
+    metaclass: Type[Any] = DeclarativeMeta,
 ) -> Any:
     r"""Construct a base class for declarative class definitions.
 
@@ -736,8 +787,19 @@ class registry:
 
     """
 
+    _class_registry: clsregistry._ClsRegistryType
+    _managers: weakref.WeakKeyDictionary[ClassManager[Any], Literal[True]]
+    _non_primary_mappers: weakref.WeakKeyDictionary[Mapper[Any], Literal[True]]
+    metadata: MetaData
+    constructor: CallableReference[Callable[..., None]]
+    type_annotation_map: _MutableTypeAnnotationMapType
+    _dependents: Set[_RegistryType]
+    _dependencies: Set[_RegistryType]
+    _new_mappers: bool
+
     def __init__(
         self,
+        *,
         metadata: Optional[MetaData] = None,
         class_registry: Optional[clsregistry._ClsRegistryType] = None,
         type_annotation_map: Optional[_TypeAnnotationMapType] = None,
@@ -799,9 +861,7 @@ class registry:
 
     def update_type_annotation_map(
         self,
-        type_annotation_map: Mapping[
-            Type, Union[Type[TypeEngine], TypeEngine]
-        ],
+        type_annotation_map: _TypeAnnotationMapType,
     ) -> None:
         """update the :paramref:`_orm.registry.type_annotation_map` with new
         values."""
@@ -817,20 +877,20 @@ class registry:
         )
 
     @property
-    def mappers(self):
+    def mappers(self) -> FrozenSet[Mapper[Any]]:
         """read only collection of all :class:`_orm.Mapper` objects."""
 
         return frozenset(manager.mapper for manager in self._managers).union(
             self._non_primary_mappers
         )
 
-    def _set_depends_on(self, registry):
+    def _set_depends_on(self, registry: RegistryType) -> None:
         if registry is self:
             return
         registry._dependents.add(self)
         self._dependencies.add(registry)
 
-    def _flag_new_mapper(self, mapper):
+    def _flag_new_mapper(self, mapper: Mapper[Any]) -> None:
         mapper._ready_for_configure = True
         if self._new_mappers:
             return
@@ -839,7 +899,9 @@ class registry:
             reg._new_mappers = True
 
     @classmethod
-    def _recurse_with_dependents(cls, registries):
+    def _recurse_with_dependents(
+        cls, registries: Set[RegistryType]
+    ) -> Iterator[RegistryType]:
         todo = registries
         done = set()
         while todo:
@@ -856,7 +918,9 @@ class registry:
             todo.update(reg._dependents.difference(done))
 
     @classmethod
-    def _recurse_with_dependencies(cls, registries):
+    def _recurse_with_dependencies(
+        cls, registries: Set[RegistryType]
+    ) -> Iterator[RegistryType]:
         todo = registries
         done = set()
         while todo:
@@ -873,7 +937,7 @@ class registry:
             # them before
             todo.update(reg._dependencies.difference(done))
 
-    def _mappers_to_configure(self):
+    def _mappers_to_configure(self) -> Iterator[Mapper[Any]]:
         return itertools.chain(
             (
                 manager.mapper
@@ -889,13 +953,13 @@ class registry:
             ),
         )
 
-    def _add_non_primary_mapper(self, np_mapper):
+    def _add_non_primary_mapper(self, np_mapper: Mapper[Any]) -> None:
         self._non_primary_mappers[np_mapper] = True
 
-    def _dispose_cls(self, cls):
+    def _dispose_cls(self, cls: Type[_O]) -> None:
         clsregistry.remove_class(cls.__name__, cls, self._class_registry)
 
-    def _add_manager(self, manager):
+    def _add_manager(self, manager: ClassManager[Any]) -> None:
         self._managers[manager] = True
         if manager.is_mapped:
             raise exc.ArgumentError(
@@ -905,7 +969,7 @@ class registry:
         assert manager.registry is None
         manager.registry = self
 
-    def configure(self, cascade=False):
+    def configure(self, cascade: bool = False) -> None:
         """Configure all as-yet unconfigured mappers in this
         :class:`_orm.registry`.
 
@@ -946,7 +1010,7 @@ class registry:
         """
         mapperlib._configure_registries({self}, cascade=cascade)
 
-    def dispose(self, cascade=False):
+    def dispose(self, cascade: bool = False) -> None:
         """Dispose of all mappers in this :class:`_orm.registry`.
 
         After invocation, all the classes that were mapped within this registry
@@ -972,7 +1036,7 @@ class registry:
 
         mapperlib._dispose_registries({self}, cascade=cascade)
 
-    def _dispose_manager_and_mapper(self, manager):
+    def _dispose_manager_and_mapper(self, manager: ClassManager[Any]) -> None:
         if "mapper" in manager.__dict__:
             mapper = manager.mapper
 
@@ -984,11 +1048,11 @@ class registry:
 
     def generate_base(
         self,
-        mapper=None,
-        cls=object,
-        name="Base",
-        metaclass=DeclarativeMeta,
-    ):
+        mapper: Optional[Callable[..., Mapper[Any]]] = None,
+        cls: Type[Any] = object,
+        name: str = "Base",
+        metaclass: Type[Any] = DeclarativeMeta,
+    ) -> Any:
         """Generate a declarative base class.
 
         Classes that inherit from the returned class object will be
@@ -1070,7 +1134,7 @@ class registry:
 
         if hasattr(cls, "__class_getitem__"):
 
-            def __class_getitem__(cls, key):
+            def __class_getitem__(cls: Type[_T], key: str) -> Type[_T]:
                 # allow generic classes in py3.9+
                 return cls
 
@@ -1078,7 +1142,7 @@ class registry:
 
         return metaclass(name, bases, class_dict)
 
-    def mapped(self, cls):
+    def mapped(self, cls: Type[_O]) -> Type[_O]:
         """Class decorator that will apply the Declarative mapping process
         to a given class.
 
@@ -1114,7 +1178,7 @@ class registry:
         _as_declarative(self, cls, cls.__dict__)
         return cls
 
-    def as_declarative_base(self, **kw):
+    def as_declarative_base(self, **kw: Any) -> Callable[[Type[_T]], Type[_T]]:
         """
         Class decorator which will invoke
         :meth:`_orm.registry.generate_base`
@@ -1142,14 +1206,14 @@ class registry:
 
         """
 
-        def decorate(cls):
+        def decorate(cls: Type[_T]) -> Type[_T]:
             kw["cls"] = cls
             kw["name"] = cls.__name__
-            return self.generate_base(**kw)
+            return self.generate_base(**kw)  # type: ignore
 
         return decorate
 
-    def map_declaratively(self, cls):
+    def map_declaratively(self, cls: Type[_O]) -> Mapper[_O]:
         """Map a class declaratively.
 
         In this form of mapping, the class is scanned for mapping information,
@@ -1194,9 +1258,15 @@ class registry:
             :meth:`_orm.registry.map_imperatively`
 
         """
-        return _as_declarative(self, cls, cls.__dict__)
+        _as_declarative(self, cls, cls.__dict__)
+        return cls.__mapper__  # type: ignore
 
-    def map_imperatively(self, class_, local_table=None, **kw):
+    def map_imperatively(
+        self,
+        class_: Type[_O],
+        local_table: Optional[FromClause] = None,
+        **kw: Any,
+    ) -> Mapper[_O]:
         r"""Map a class imperatively.
 
         In this form of mapping, the class is not scanned for any mapping
@@ -1251,7 +1321,7 @@ class registry:
 RegistryType = registry
 
 
-def as_declarative(**kw):
+def as_declarative(**kw: Any) -> Callable[[Type[_T]], Type[_T]]:
     """
     Class decorator which will adapt a given class into a
     :func:`_orm.declarative_base`.
@@ -1292,14 +1362,9 @@ def as_declarative(**kw):
 @inspection._inspects(
     DeclarativeMeta, DeclarativeBase, DeclarativeAttributeIntercept
 )
-def _inspect_decl_meta(cls: Type[Any]) -> Mapper[Any]:
-    mp: Mapper[Any] = _inspect_mapped_class(cls)
+def _inspect_decl_meta(cls: Type[Any]) -> Optional[Mapper[Any]]:
+    mp: Optional[Mapper[Any]] = _inspect_mapped_class(cls)
     if mp is None:
         if _DeferredMapperConfig.has_cls(cls):
             _DeferredMapperConfig.raise_unmapped_for_cls(cls)
-            raise orm_exc.UnmappedClassError(
-                cls,
-                msg="Class %s has a deferred mapping on it.  It is not yet "
-                "usable as a mapped class." % orm_exc._safe_cls_name(cls),
-            )
     return mp