plugin.py 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284
  1. # ext/mypy/plugin.py
  2. # Copyright (C) 2021 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: https://www.opensource.org/licenses/mit-license.php
  7. """
  8. Mypy plugin for SQLAlchemy ORM.
  9. """
  10. from typing import Callable
  11. from typing import List
  12. from typing import Optional
  13. from typing import Tuple
  14. from typing import Type as TypingType
  15. from typing import Union
  16. from mypy import nodes
  17. from mypy.mro import calculate_mro
  18. from mypy.mro import MroError
  19. from mypy.nodes import Block
  20. from mypy.nodes import ClassDef
  21. from mypy.nodes import GDEF
  22. from mypy.nodes import MypyFile
  23. from mypy.nodes import NameExpr
  24. from mypy.nodes import SymbolTable
  25. from mypy.nodes import SymbolTableNode
  26. from mypy.nodes import TypeInfo
  27. from mypy.plugin import AttributeContext
  28. from mypy.plugin import ClassDefContext
  29. from mypy.plugin import DynamicClassDefContext
  30. from mypy.plugin import Plugin
  31. from mypy.plugin import SemanticAnalyzerPluginInterface
  32. from mypy.types import get_proper_type
  33. from mypy.types import Instance
  34. from mypy.types import Type
  35. from . import decl_class
  36. from . import names
  37. from . import util
  38. class SQLAlchemyPlugin(Plugin):
  39. def get_dynamic_class_hook(
  40. self, fullname: str
  41. ) -> Optional[Callable[[DynamicClassDefContext], None]]:
  42. if names.type_id_for_fullname(fullname) is names.DECLARATIVE_BASE:
  43. return _dynamic_class_hook
  44. return None
  45. def get_customize_class_mro_hook(
  46. self, fullname: str
  47. ) -> Optional[Callable[[ClassDefContext], None]]:
  48. return _fill_in_decorators
  49. def get_class_decorator_hook(
  50. self, fullname: str
  51. ) -> Optional[Callable[[ClassDefContext], None]]:
  52. sym = self.lookup_fully_qualified(fullname)
  53. if sym is not None and sym.node is not None:
  54. type_id = names.type_id_for_named_node(sym.node)
  55. if type_id is names.MAPPED_DECORATOR:
  56. return _cls_decorator_hook
  57. elif type_id in (
  58. names.AS_DECLARATIVE,
  59. names.AS_DECLARATIVE_BASE,
  60. ):
  61. return _base_cls_decorator_hook
  62. elif type_id is names.DECLARATIVE_MIXIN:
  63. return _declarative_mixin_hook
  64. return None
  65. def get_metaclass_hook(
  66. self, fullname: str
  67. ) -> Optional[Callable[[ClassDefContext], None]]:
  68. if names.type_id_for_fullname(fullname) is names.DECLARATIVE_META:
  69. # Set any classes that explicitly have metaclass=DeclarativeMeta
  70. # as declarative so the check in `get_base_class_hook()` works
  71. return _metaclass_cls_hook
  72. return None
  73. def get_base_class_hook(
  74. self, fullname: str
  75. ) -> Optional[Callable[[ClassDefContext], None]]:
  76. sym = self.lookup_fully_qualified(fullname)
  77. if (
  78. sym
  79. and isinstance(sym.node, TypeInfo)
  80. and util.has_declarative_base(sym.node)
  81. ):
  82. return _base_cls_hook
  83. return None
  84. def get_attribute_hook(
  85. self, fullname: str
  86. ) -> Optional[Callable[[AttributeContext], Type]]:
  87. if fullname.startswith(
  88. "sqlalchemy.orm.attributes.QueryableAttribute."
  89. ):
  90. return _queryable_getattr_hook
  91. return None
  92. def get_additional_deps(
  93. self, file: MypyFile
  94. ) -> List[Tuple[int, str, int]]:
  95. return [
  96. (10, "sqlalchemy.orm.attributes", -1),
  97. (10, "sqlalchemy.orm.decl_api", -1),
  98. ]
  99. def plugin(version: str) -> TypingType[SQLAlchemyPlugin]:
  100. return SQLAlchemyPlugin
  101. def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None:
  102. """Generate a declarative Base class when the declarative_base() function
  103. is encountered."""
  104. _add_globals(ctx)
  105. cls = ClassDef(ctx.name, Block([]))
  106. cls.fullname = ctx.api.qualified_name(ctx.name)
  107. info = TypeInfo(SymbolTable(), cls, ctx.api.cur_mod_id)
  108. cls.info = info
  109. _set_declarative_metaclass(ctx.api, cls)
  110. cls_arg = util.get_callexpr_kwarg(ctx.call, "cls", expr_types=(NameExpr,))
  111. if cls_arg is not None and isinstance(cls_arg.node, TypeInfo):
  112. util.set_is_base(cls_arg.node)
  113. decl_class.scan_declarative_assignments_and_apply_types(
  114. cls_arg.node.defn, ctx.api, is_mixin_scan=True
  115. )
  116. info.bases = [Instance(cls_arg.node, [])]
  117. else:
  118. obj = ctx.api.named_type(names.NAMED_TYPE_BUILTINS_OBJECT)
  119. info.bases = [obj]
  120. try:
  121. calculate_mro(info)
  122. except MroError:
  123. util.fail(
  124. ctx.api, "Not able to calculate MRO for declarative base", ctx.call
  125. )
  126. obj = ctx.api.named_type(names.NAMED_TYPE_BUILTINS_OBJECT)
  127. info.bases = [obj]
  128. info.fallback_to_any = True
  129. ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
  130. util.set_is_base(info)
  131. def _fill_in_decorators(ctx: ClassDefContext) -> None:
  132. for decorator in ctx.cls.decorators:
  133. # set the ".fullname" attribute of a class decorator
  134. # that is a MemberExpr. This causes the logic in
  135. # semanal.py->apply_class_plugin_hooks to invoke the
  136. # get_class_decorator_hook for our "registry.map_class()"
  137. # and "registry.as_declarative_base()" methods.
  138. # this seems like a bug in mypy that these decorators are otherwise
  139. # skipped.
  140. if (
  141. isinstance(decorator, nodes.CallExpr)
  142. and isinstance(decorator.callee, nodes.MemberExpr)
  143. and decorator.callee.name == "as_declarative_base"
  144. ):
  145. target = decorator.callee
  146. elif (
  147. isinstance(decorator, nodes.MemberExpr)
  148. and decorator.name == "mapped"
  149. ):
  150. target = decorator
  151. else:
  152. continue
  153. assert isinstance(target.expr, NameExpr)
  154. sym = ctx.api.lookup_qualified(
  155. target.expr.name, target, suppress_errors=True
  156. )
  157. if sym and sym.node:
  158. sym_type = get_proper_type(sym.type)
  159. if isinstance(sym_type, Instance):
  160. target.fullname = f"{sym_type.type.fullname}.{target.name}"
  161. else:
  162. # if the registry is in the same file as where the
  163. # decorator is used, it might not have semantic
  164. # symbols applied and we can't get a fully qualified
  165. # name or an inferred type, so we are actually going to
  166. # flag an error in this case that they need to annotate
  167. # it. The "registry" is declared just
  168. # once (or few times), so they have to just not use
  169. # type inference for its assignment in this one case.
  170. util.fail(
  171. ctx.api,
  172. "Class decorator called %s(), but we can't "
  173. "tell if it's from an ORM registry. Please "
  174. "annotate the registry assignment, e.g. "
  175. "my_registry: registry = registry()" % target.name,
  176. sym.node,
  177. )
  178. def _cls_decorator_hook(ctx: ClassDefContext) -> None:
  179. _add_globals(ctx)
  180. assert isinstance(ctx.reason, nodes.MemberExpr)
  181. expr = ctx.reason.expr
  182. assert isinstance(expr, nodes.RefExpr) and isinstance(expr.node, nodes.Var)
  183. node_type = get_proper_type(expr.node.type)
  184. assert (
  185. isinstance(node_type, Instance)
  186. and names.type_id_for_named_node(node_type.type) is names.REGISTRY
  187. )
  188. decl_class.scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
  189. def _base_cls_decorator_hook(ctx: ClassDefContext) -> None:
  190. _add_globals(ctx)
  191. cls = ctx.cls
  192. _set_declarative_metaclass(ctx.api, cls)
  193. util.set_is_base(ctx.cls.info)
  194. decl_class.scan_declarative_assignments_and_apply_types(
  195. cls, ctx.api, is_mixin_scan=True
  196. )
  197. def _declarative_mixin_hook(ctx: ClassDefContext) -> None:
  198. _add_globals(ctx)
  199. util.set_is_base(ctx.cls.info)
  200. decl_class.scan_declarative_assignments_and_apply_types(
  201. ctx.cls, ctx.api, is_mixin_scan=True
  202. )
  203. def _metaclass_cls_hook(ctx: ClassDefContext) -> None:
  204. util.set_is_base(ctx.cls.info)
  205. def _base_cls_hook(ctx: ClassDefContext) -> None:
  206. _add_globals(ctx)
  207. decl_class.scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
  208. def _queryable_getattr_hook(ctx: AttributeContext) -> Type:
  209. # how do I....tell it it has no attribute of a certain name?
  210. # can't find any Type that seems to match that
  211. return ctx.default_attr_type
  212. def _add_globals(ctx: Union[ClassDefContext, DynamicClassDefContext]) -> None:
  213. """Add __sa_DeclarativeMeta and __sa_Mapped symbol to the global space
  214. for all class defs
  215. """
  216. util.add_global(ctx, "sqlalchemy.orm.attributes", "Mapped", "__sa_Mapped")
  217. def _set_declarative_metaclass(
  218. api: SemanticAnalyzerPluginInterface, target_cls: ClassDef
  219. ) -> None:
  220. info = target_cls.info
  221. sym = api.lookup_fully_qualified_or_none(
  222. "sqlalchemy.orm.decl_api.DeclarativeMeta"
  223. )
  224. assert sym is not None and isinstance(sym.node, TypeInfo)
  225. info.declared_metaclass = info.metaclass_type = Instance(sym.node, [])