From 974f414ea41bd8a28fabbf3cd5786e1532eeb91d Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Tue, 23 Mar 2021 18:49:28 -0400 Subject: [PATCH] support as_declarative, as_declarative_base Added support for the Mypy extension to correctly interpret a declarative base class that's generated using the :func:`_orm.as_declarative` function as well as the :meth:`_orm.registry.as_declarative_base` method. Change-Id: I227f4abebe157a7df3f8772893bbea6669cc8555 --- .../changelog/unreleased_14/asdecl_mypy.rst | 6 + lib/sqlalchemy/ext/mypy/names.py | 18 ++- lib/sqlalchemy/ext/mypy/plugin.py | 109 ++++++++++++------ test/ext/mypy/files/as_declarative.py | 26 +++++ test/ext/mypy/files/as_declarative_base.py | 28 +++++ 5 files changed, 148 insertions(+), 39 deletions(-) create mode 100644 doc/build/changelog/unreleased_14/asdecl_mypy.rst create mode 100644 test/ext/mypy/files/as_declarative.py create mode 100644 test/ext/mypy/files/as_declarative_base.py diff --git a/doc/build/changelog/unreleased_14/asdecl_mypy.rst b/doc/build/changelog/unreleased_14/asdecl_mypy.rst new file mode 100644 index 0000000000..170b39f5e8 --- /dev/null +++ b/doc/build/changelog/unreleased_14/asdecl_mypy.rst @@ -0,0 +1,6 @@ +.. change:: + :tags: bug, mypy + + Added support for the Mypy extension to correctly interpret a declarative + base class that's generated using the :func:`_orm.as_declarative` function + as well as the :meth:`_orm.registry.as_declarative_base` method. diff --git a/lib/sqlalchemy/ext/mypy/names.py b/lib/sqlalchemy/ext/mypy/names.py index c9d48fcd82..d1fd77415a 100644 --- a/lib/sqlalchemy/ext/mypy/names.py +++ b/lib/sqlalchemy/ext/mypy/names.py @@ -34,7 +34,8 @@ SYNONYM_PROPERTY = util.symbol("SYNONYM_PROPERTY") COMPOSITE_PROPERTY = util.symbol("COMPOSITE_PROPERTY") DECLARED_ATTR = util.symbol("DECLARED_ATTR") MAPPER_PROPERTY = util.symbol("MAPPER_PROPERTY") - +AS_DECLARATIVE = util.symbol("AS_DECLARATIVE") +AS_DECLARATIVE_BASE = util.symbol("AS_DECLARATIVE_BASE") _lookup = { "Column": ( @@ -111,6 +112,21 @@ _lookup = { "sqlalchemy.orm.registry.mapped", }, ), + "as_declarative": ( + AS_DECLARATIVE, + { + "sqlalchemy.ext.declarative.as_declarative", + "sqlalchemy.orm.decl_api.as_declarative", + "sqlalchemy.orm.as_declarative", + }, + ), + "as_declarative_base": ( + AS_DECLARATIVE_BASE, + { + "sqlalchemy.orm.decl_api.registry.as_declarative_base", + "sqlalchemy.orm.registry.as_declarative_base", + }, + ), "declared_attr": ( DECLARED_ATTR, { diff --git a/lib/sqlalchemy/ext/mypy/plugin.py b/lib/sqlalchemy/ext/mypy/plugin.py index 9fcd09b1ee..c8fbcd6a21 100644 --- a/lib/sqlalchemy/ext/mypy/plugin.py +++ b/lib/sqlalchemy/ext/mypy/plugin.py @@ -30,6 +30,7 @@ from mypy.plugin import ClassDefContext from mypy.plugin import DynamicClassDefContext from mypy.plugin import Optional from mypy.plugin import Plugin +from mypy.plugin import SemanticAnalyzerPluginInterface from mypy.types import Instance from . import decl_class @@ -69,13 +70,18 @@ class CustomPlugin(Plugin): ) -> Optional[Callable[[ClassDefContext], None]]: sym = self.lookup_fully_qualified(fullname) - if ( sym is not None and names._type_id_for_named_node(sym.node) is names.MAPPED_DECORATOR ): return _cls_decorator_hook + elif sym is not None and names._type_id_for_named_node(sym.node) in ( + names.AS_DECLARATIVE, + names.AS_DECLARATIVE_BASE, + ): + return _base_cls_decorator_hook + return None def get_customize_class_mro_hook( @@ -116,39 +122,47 @@ def _fill_in_decorators(ctx: ClassDefContext) -> None: # set the ".fullname" attribute of a class decorator # that is a MemberExpr. This causes the logic in # semanal.py->apply_class_plugin_hooks to invoke the - # get_class_decorator_hook for our "registry.map_class()" method. + # get_class_decorator_hook for our "registry.map_class()" + # and "registry.as_declarative_base()" methods. # this seems like a bug in mypy that these decorators are otherwise # skipped. if ( + isinstance(decorator, nodes.CallExpr) + and isinstance(decorator.callee, nodes.MemberExpr) + and decorator.callee.name == "as_declarative_base" + ): + target = decorator.callee + elif ( isinstance(decorator, nodes.MemberExpr) and decorator.name == "mapped" ): - - sym = ctx.api.lookup( - decorator.expr.name, decorator, suppress_errors=True - ) - if sym: - if sym.node.type and hasattr(sym.node.type, "type"): - decorator.fullname = ( - f"{sym.node.type.type.fullname}.{decorator.name}" - ) - else: - # if the registry is in the same file as where the - # decorator is used, it might not have semantic - # symbols applied and we can't get a fully qualified - # name or an inferred type, so we are actually going to - # flag an error in this case that they need to annotate - # it. The "registry" is declared just - # once (or few times), so they have to just not use - # type inference for its assignment in this one case. - util.fail( - ctx.api, - "Class decorator called mapped(), but we can't " - "tell if it's from an ORM registry. Please " - "annotate the registry assignment, e.g. " - "my_registry: registry = registry()", - sym.node, - ) + target = decorator + else: + continue + + sym = ctx.api.lookup(target.expr.name, target, suppress_errors=True) + if sym: + if sym.node.type and hasattr(sym.node.type, "type"): + target.fullname = ( + f"{sym.node.type.type.fullname}.{target.name}" + ) + else: + # if the registry is in the same file as where the + # decorator is used, it might not have semantic + # symbols applied and we can't get a fully qualified + # name or an inferred type, so we are actually going to + # flag an error in this case that they need to annotate + # it. The "registry" is declared just + # once (or few times), so they have to just not use + # type inference for its assignment in this one case. + util.fail( + ctx.api, + "Class decorator called %s(), but we can't " + "tell if it's from an ORM registry. Please " + "annotate the registry assignment, e.g. " + "my_registry: registry = registry()" % target.name, + sym.node, + ) def _cls_metadata_hook(ctx: ClassDefContext) -> None: @@ -167,14 +181,10 @@ def _cls_decorator_hook(ctx: ClassDefContext) -> None: decl_class._scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api) -def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None: - """Generate a declarative Base class when the declarative_base() function - is encountered.""" - - cls = ClassDef(ctx.name, Block([])) - cls.fullname = ctx.api.qualified_name(ctx.name) - - declarative_meta_sym: SymbolTableNode = ctx.api.modules[ +def _make_declarative_meta( + api: SemanticAnalyzerPluginInterface, target_cls: ClassDef +): + declarative_meta_sym: SymbolTableNode = api.modules[ "sqlalchemy.orm.decl_api" ].names["DeclarativeMeta"] declarative_meta_typeinfo: TypeInfo = declarative_meta_sym.node @@ -184,13 +194,35 @@ def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None: declarative_meta_name.fullname = "sqlalchemy.orm.decl_api.DeclarativeMeta" declarative_meta_name.node = declarative_meta_typeinfo - cls.metaclass = declarative_meta_name + target_cls.metaclass = declarative_meta_name declarative_meta_instance = Instance(declarative_meta_typeinfo, []) - info = TypeInfo(SymbolTable(), cls, ctx.api.cur_mod_id) + info = target_cls.info info.declared_metaclass = info.metaclass_type = declarative_meta_instance + + +def _base_cls_decorator_hook(ctx: ClassDefContext) -> None: + + cls = ctx.cls + + _make_declarative_meta(ctx.api, cls) + + decl_class._scan_declarative_assignments_and_apply_types( + cls, ctx.api, is_mixin_scan=True + ) + + +def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None: + """Generate a declarative Base class when the declarative_base() function + is encountered.""" + + cls = ClassDef(ctx.name, Block([])) + cls.fullname = ctx.api.qualified_name(ctx.name) + + info = TypeInfo(SymbolTable(), cls, ctx.api.cur_mod_id) cls.info = info + _make_declarative_meta(ctx.api, cls) cls_arg = util._get_callexpr_kwarg(ctx.call, "cls") if cls_arg is not None: @@ -209,6 +241,7 @@ def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None: util.fail( ctx.api, "Not able to calculate MRO for declarative base", ctx.call ) + obj = ctx.api.builtin_type("builtins.object") info.bases = [obj] info.fallback_to_any = True diff --git a/test/ext/mypy/files/as_declarative.py b/test/ext/mypy/files/as_declarative.py new file mode 100644 index 0000000000..7a3fdc068d --- /dev/null +++ b/test/ext/mypy/files/as_declarative.py @@ -0,0 +1,26 @@ +from sqlalchemy import Column +from sqlalchemy import Integer +from sqlalchemy import String +from sqlalchemy.ext.declarative import as_declarative + + +@as_declarative +class Base(object): + updated_at = Column(Integer) + + +class Foo(Base): + __tablename__ = "foo" + id: int = Column(Integer(), primary_key=True) + name: str = Column(String) + + +f1 = Foo() + +val: int = f1.id + +p: str = f1.name + +Foo.id.property + +f2 = Foo(name="some name", updated_at=5) diff --git a/test/ext/mypy/files/as_declarative_base.py b/test/ext/mypy/files/as_declarative_base.py new file mode 100644 index 0000000000..efba8c362a --- /dev/null +++ b/test/ext/mypy/files/as_declarative_base.py @@ -0,0 +1,28 @@ +from sqlalchemy import Column +from sqlalchemy import Integer +from sqlalchemy import String +from sqlalchemy.orm import registry + +reg: registry = registry() + + +@reg.as_declarative_base() +class Base(object): + updated_at = Column(Integer) + + +class Foo(Base): + __tablename__ = "foo" + id: int = Column(Integer(), primary_key=True) + name: str = Column(String) + + +f1 = Foo() + +val: int = f1.id + +p: str = f1.name + +Foo.id.property + +f2 = Foo(name="some name", updated_at=5) -- 2.47.2