123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284 |
- # ext/mypy/plugin.py
- # Copyright (C) 2021 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
- """
- Mypy plugin for SQLAlchemy ORM.
- """
- from typing import Callable
- from typing import List
- from typing import Optional
- from typing import Tuple
- from typing import Type as TypingType
- from typing import Union
- from mypy import nodes
- from mypy.mro import calculate_mro
- from mypy.mro import MroError
- from mypy.nodes import Block
- from mypy.nodes import ClassDef
- from mypy.nodes import GDEF
- from mypy.nodes import MypyFile
- from mypy.nodes import NameExpr
- from mypy.nodes import SymbolTable
- from mypy.nodes import SymbolTableNode
- from mypy.nodes import TypeInfo
- from mypy.plugin import AttributeContext
- from mypy.plugin import ClassDefContext
- from mypy.plugin import DynamicClassDefContext
- from mypy.plugin import Plugin
- from mypy.plugin import SemanticAnalyzerPluginInterface
- from mypy.types import get_proper_type
- from mypy.types import Instance
- from mypy.types import Type
- from . import decl_class
- from . import names
- from . import util
- class SQLAlchemyPlugin(Plugin):
- def get_dynamic_class_hook(
- self, fullname: str
- ) -> Optional[Callable[[DynamicClassDefContext], None]]:
- if names.type_id_for_fullname(fullname) is names.DECLARATIVE_BASE:
- return _dynamic_class_hook
- return None
- def get_customize_class_mro_hook(
- self, fullname: str
- ) -> Optional[Callable[[ClassDefContext], None]]:
- return _fill_in_decorators
- def get_class_decorator_hook(
- self, fullname: str
- ) -> Optional[Callable[[ClassDefContext], None]]:
- sym = self.lookup_fully_qualified(fullname)
- if sym is not None and sym.node is not None:
- type_id = names.type_id_for_named_node(sym.node)
- if type_id is names.MAPPED_DECORATOR:
- return _cls_decorator_hook
- elif type_id in (
- names.AS_DECLARATIVE,
- names.AS_DECLARATIVE_BASE,
- ):
- return _base_cls_decorator_hook
- elif type_id is names.DECLARATIVE_MIXIN:
- return _declarative_mixin_hook
- return None
- def get_metaclass_hook(
- self, fullname: str
- ) -> Optional[Callable[[ClassDefContext], None]]:
- if names.type_id_for_fullname(fullname) is names.DECLARATIVE_META:
- # Set any classes that explicitly have metaclass=DeclarativeMeta
- # as declarative so the check in `get_base_class_hook()` works
- return _metaclass_cls_hook
- return None
- def get_base_class_hook(
- self, fullname: str
- ) -> Optional[Callable[[ClassDefContext], None]]:
- sym = self.lookup_fully_qualified(fullname)
- if (
- sym
- and isinstance(sym.node, TypeInfo)
- and util.has_declarative_base(sym.node)
- ):
- return _base_cls_hook
- return None
- def get_attribute_hook(
- self, fullname: str
- ) -> Optional[Callable[[AttributeContext], Type]]:
- if fullname.startswith(
- "sqlalchemy.orm.attributes.QueryableAttribute."
- ):
- return _queryable_getattr_hook
- return None
- def get_additional_deps(
- self, file: MypyFile
- ) -> List[Tuple[int, str, int]]:
- return [
- (10, "sqlalchemy.orm.attributes", -1),
- (10, "sqlalchemy.orm.decl_api", -1),
- ]
- def plugin(version: str) -> TypingType[SQLAlchemyPlugin]:
- return SQLAlchemyPlugin
- def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None:
- """Generate a declarative Base class when the declarative_base() function
- is encountered."""
- _add_globals(ctx)
- 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
- _set_declarative_metaclass(ctx.api, cls)
- cls_arg = util.get_callexpr_kwarg(ctx.call, "cls", expr_types=(NameExpr,))
- if cls_arg is not None and isinstance(cls_arg.node, TypeInfo):
- util.set_is_base(cls_arg.node)
- decl_class.scan_declarative_assignments_and_apply_types(
- cls_arg.node.defn, ctx.api, is_mixin_scan=True
- )
- info.bases = [Instance(cls_arg.node, [])]
- else:
- obj = ctx.api.named_type(names.NAMED_TYPE_BUILTINS_OBJECT)
- info.bases = [obj]
- try:
- calculate_mro(info)
- except MroError:
- util.fail(
- ctx.api, "Not able to calculate MRO for declarative base", ctx.call
- )
- obj = ctx.api.named_type(names.NAMED_TYPE_BUILTINS_OBJECT)
- info.bases = [obj]
- info.fallback_to_any = True
- ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
- util.set_is_base(info)
- def _fill_in_decorators(ctx: ClassDefContext) -> None:
- for decorator in ctx.cls.decorators:
- # 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()"
- # 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"
- ):
- target = decorator
- else:
- continue
- assert isinstance(target.expr, NameExpr)
- sym = ctx.api.lookup_qualified(
- target.expr.name, target, suppress_errors=True
- )
- if sym and sym.node:
- sym_type = get_proper_type(sym.type)
- if isinstance(sym_type, Instance):
- target.fullname = f"{sym_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_decorator_hook(ctx: ClassDefContext) -> None:
- _add_globals(ctx)
- assert isinstance(ctx.reason, nodes.MemberExpr)
- expr = ctx.reason.expr
- assert isinstance(expr, nodes.RefExpr) and isinstance(expr.node, nodes.Var)
- node_type = get_proper_type(expr.node.type)
- assert (
- isinstance(node_type, Instance)
- and names.type_id_for_named_node(node_type.type) is names.REGISTRY
- )
- decl_class.scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
- def _base_cls_decorator_hook(ctx: ClassDefContext) -> None:
- _add_globals(ctx)
- cls = ctx.cls
- _set_declarative_metaclass(ctx.api, cls)
- util.set_is_base(ctx.cls.info)
- decl_class.scan_declarative_assignments_and_apply_types(
- cls, ctx.api, is_mixin_scan=True
- )
- def _declarative_mixin_hook(ctx: ClassDefContext) -> None:
- _add_globals(ctx)
- util.set_is_base(ctx.cls.info)
- decl_class.scan_declarative_assignments_and_apply_types(
- ctx.cls, ctx.api, is_mixin_scan=True
- )
- def _metaclass_cls_hook(ctx: ClassDefContext) -> None:
- util.set_is_base(ctx.cls.info)
- def _base_cls_hook(ctx: ClassDefContext) -> None:
- _add_globals(ctx)
- decl_class.scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
- def _queryable_getattr_hook(ctx: AttributeContext) -> Type:
- # how do I....tell it it has no attribute of a certain name?
- # can't find any Type that seems to match that
- return ctx.default_attr_type
- def _add_globals(ctx: Union[ClassDefContext, DynamicClassDefContext]) -> None:
- """Add __sa_DeclarativeMeta and __sa_Mapped symbol to the global space
- for all class defs
- """
- util.add_global(ctx, "sqlalchemy.orm.attributes", "Mapped", "__sa_Mapped")
- def _set_declarative_metaclass(
- api: SemanticAnalyzerPluginInterface, target_cls: ClassDef
- ) -> None:
- info = target_cls.info
- sym = api.lookup_fully_qualified_or_none(
- "sqlalchemy.orm.decl_api.DeclarativeMeta"
- )
- assert sym is not None and isinstance(sym.node, TypeInfo)
- info.declared_metaclass = info.metaclass_type = Instance(sym.node, [])
|