From ef17852999a4c0e01a84a83da52b18f583a3de42 Mon Sep 17 00:00:00 2001
From: Sebastian Rittau <srittau@rittau.biz>
Date: Tue, 21 Dec 2021 16:40:37 +0100
Subject: [PATCH] SQLAlchemy: Partly annotate declarative API (#6647)

---
 stubs/SQLAlchemy/sqlalchemy/orm/decl_api.pyi | 117 +++++++++++++++----
 1 file changed, 95 insertions(+), 22 deletions(-)

diff --git a/stubs/SQLAlchemy/sqlalchemy/orm/decl_api.pyi b/stubs/SQLAlchemy/sqlalchemy/orm/decl_api.pyi
index 4988743b2..c62611662 100644
--- a/stubs/SQLAlchemy/sqlalchemy/orm/decl_api.pyi
+++ b/stubs/SQLAlchemy/sqlalchemy/orm/decl_api.pyi
@@ -1,19 +1,38 @@
-from typing import Any
+from collections.abc import Callable
+from typing import Any, ClassVar, Type, TypeVar, overload
 
+from ..engine.interfaces import Connectable
+from ..sql.schema import MetaData
 from ..util import hybridproperty
 from . import interfaces
 
-def has_inherited_table(cls): ...
+_ClsT = TypeVar("_ClsT", bound=Type[Any])
+_DeclT = TypeVar("_DeclT", bound=Type[_DeclarativeBase])
+
+# Dynamic class as created by registry.generate_base() via DeclarativeMeta
+# or another metaclass. This class does not exist at runtime.
+class _DeclarativeBase(Any):  # super classes are dynamic
+    registry: ClassVar[registry]
+    metadata: ClassVar[MetaData]
+    __abstract__: ClassVar[bool]
+    # not always existing:
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+    __mapper_cls__: ClassVar[Any]
+    __class_getitem__: ClassVar[Any]
+
+# Meta class (or function) that creates a _DeclarativeBase class.
+_DeclarativeBaseMeta = Callable[[str, tuple[Type[Any], ...], dict[str, Any]], _DeclT]
+
+def has_inherited_table(cls: Type[Any]) -> bool: ...
 
 class DeclarativeMeta(type):
-    def __init__(cls, classname, bases, dict_, **kw) -> None: ...
-    def __setattr__(cls, key, value) -> None: ...
-    def __delattr__(cls, key) -> None: ...
+    def __init__(cls, classname: str, bases: tuple[Type[Any], ...], dict_: dict[str, Any], **kw: object) -> None: ...
+    def __setattr__(cls, key: str, value: Any) -> None: ...
+    def __delattr__(cls, key: str) -> None: ...
 
 def synonym_for(name, map_column: bool = ...): ...
 
 class declared_attr(interfaces._MappedAttribute, property):
-    __doc__: Any
     def __init__(self, fget, cascading: bool = ...) -> None: ...
     def __get__(self, self_, cls): ...
     @hybridproperty
@@ -24,32 +43,86 @@ class _stateful_declared_attr(declared_attr):
     def __init__(self, **kw) -> None: ...
     def __call__(self, fn): ...
 
-def declarative_mixin(cls): ...
+def declarative_mixin(cls: _ClsT) -> _ClsT: ...
+@overload
 def declarative_base(
-    bind: Any | None = ...,
-    metadata: Any | None = ...,
+    bind: Connectable | None = ...,
+    metadata: MetaData | None = ...,
     mapper: Any | None = ...,
-    cls=...,
+    cls: Type[Any] | tuple[Type[Any], ...] = ...,
     name: str = ...,
-    constructor=...,
-    class_registry: Any | None = ...,
-    metaclass=...,
-): ...
+    constructor: Callable[..., None] = ...,
+    class_registry: dict[str, Type[Any]] | None = ...,
+) -> Type[_DeclarativeBase]: ...
+@overload
+def declarative_base(
+    bind: Connectable | None = ...,
+    metadata: MetaData | None = ...,
+    mapper: Any | None = ...,
+    cls: Type[Any] | tuple[Type[Any], ...] = ...,
+    name: str = ...,
+    constructor: Callable[..., None] = ...,
+    class_registry: dict[str, Type[Any]] | None = ...,
+    *,
+    metaclass: _DeclarativeBaseMeta[_DeclT],
+) -> _DeclT: ...
+@overload
+def declarative_base(
+    bind: Connectable | None,
+    metadata: MetaData | None,
+    mapper: Any | None,
+    cls: Type[Any] | tuple[Type[Any], ...],
+    name: str,
+    constructor: Callable[..., None],
+    class_registry: dict[str, Type[Any]] | None,
+    metaclass: _DeclarativeBaseMeta[_DeclT],
+) -> _DeclT: ...
 
 class registry:
-    metadata: Any
-    constructor: Any
+    metadata: MetaData
+    constructor: Callable[..., None]
     def __init__(
-        self, metadata: Any | None = ..., class_registry: Any | None = ..., constructor=..., _bind: Any | None = ...
+        self,
+        metadata: MetaData | None = ...,
+        class_registry: dict[str, Type[Any]] | None = ...,
+        constructor: Callable[..., None] = ...,
+        _bind: Connectable | None = ...,
     ) -> None: ...
     @property
-    def mappers(self): ...
+    def mappers(self) -> frozenset[Any]: ...
     def configure(self, cascade: bool = ...) -> None: ...
     def dispose(self, cascade: bool = ...) -> None: ...
-    def generate_base(self, mapper: Any | None = ..., cls=..., name: str = ..., metaclass=...): ...
-    def mapped(self, cls): ...
-    def as_declarative_base(self, **kw): ...
+    @overload
+    def generate_base(
+        self, mapper: Any | None = ..., cls: Type[Any] | tuple[Type[Any], ...] = ..., name: str = ...
+    ) -> Type[_DeclarativeBase]: ...
+    @overload
+    def generate_base(
+        self,
+        mapper: Any | None = ...,
+        cls: Type[Any] | tuple[Type[Any], ...] = ...,
+        name: str = ...,
+        *,
+        metaclass: _DeclarativeBaseMeta[_DeclT],
+    ) -> _DeclT: ...
+    @overload
+    def generate_base(
+        self, mapper: Any | None, cls: Type[Any] | tuple[Type[Any], ...], name: str, metaclass: _DeclarativeBaseMeta[_DeclT]
+    ) -> Type[_DeclarativeBase]: ...
+    def mapped(self, cls: _ClsT) -> _ClsT: ...
+    # Return type of the callable is a _DeclarativeBase class with the passed in class as base.
+    # This could be better approximated with Intersection[PassedInClass, _DeclarativeBase].
+    def as_declarative_base(
+        self, *, mapper: Any | None = ..., metaclass: _DeclarativeBaseMeta[_DeclT] = ...
+    ) -> Callable[[_ClsT], _ClsT | _DeclT | Any]: ...
     def map_declaratively(self, cls): ...
     def map_imperatively(self, class_, local_table: Any | None = ..., **kw): ...
 
-def as_declarative(**kw): ...
+def as_declarative(
+    *,
+    bind: Connectable | None = ...,
+    metadata: MetaData | None = ...,
+    class_registry: dict[str, Type[Any]] | None = ...,
+    mapper: Any | None = ...,
+    metaclass: _DeclarativeBaseMeta[_DeclT] = ...,
+) -> Callable[[_ClsT], _ClsT | _DeclT | Any]: ...