Update 2025-04-24_11:44:19
This commit is contained in:
oib
303
venv/lib/python3.11/site-packages/sqlalchemy/ext/mypy/plugin.py
Normal file
303
venv/lib/python3.11/site-packages/sqlalchemy/ext/mypy/plugin.py
Normal file
@ -0,0 +1,303 @@
|
||||
# ext/mypy/plugin.py
|
||||
# Copyright (C) 2021-2025 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 __future__ import annotations
|
||||
|
||||
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
|
||||
|
||||
try:
|
||||
__import__("sqlalchemy-stubs")
|
||||
except ImportError:
|
||||
pass
|
||||
else:
|
||||
raise ImportError(
|
||||
"The SQLAlchemy mypy plugin in SQLAlchemy "
|
||||
"2.0 does not work with sqlalchemy-stubs or "
|
||||
"sqlalchemy2-stubs installed, as well as with any other third party "
|
||||
"SQLAlchemy stubs. Please uninstall all SQLAlchemy stubs "
|
||||
"packages."
|
||||
)
|
||||
|
||||
|
||||
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", -1),
|
||||
(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
|
||||
|
||||
if isinstance(target.expr, NameExpr):
|
||||
sym = ctx.api.lookup_qualified(
|
||||
target.expr.name, target, suppress_errors=True
|
||||
)
|
||||
else:
|
||||
continue
|
||||
|
||||
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", "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, [])
|
||||
Reference in New Issue
Block a user