Загрузить файлы в «venv/Lib/site-packages/sqlalchemy/ext/mypy»

This commit is contained in:
2026-07-02 20:29:34 +00:00
parent a6b44fa5b5
commit 844c559243
2 changed files with 661 additions and 0 deletions

View File

@@ -0,0 +1,304 @@
# ext/mypy/plugin.py
# Copyright (C) 2021-2026 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, [])

View File

@@ -0,0 +1,357 @@
# ext/mypy/util.py
# Copyright (C) 2021-2026 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
from __future__ import annotations
import re
from typing import Any
from typing import Iterable
from typing import Iterator
from typing import List
from typing import Optional
from typing import overload
from typing import Tuple
from typing import Type as TypingType
from typing import TypeVar
from typing import Union
from mypy import version
from mypy.messages import format_type as _mypy_format_type
from mypy.nodes import CallExpr
from mypy.nodes import ClassDef
from mypy.nodes import CLASSDEF_NO_INFO
from mypy.nodes import Context
from mypy.nodes import Expression
from mypy.nodes import FuncDef
from mypy.nodes import IfStmt
from mypy.nodes import JsonDict
from mypy.nodes import MemberExpr
from mypy.nodes import NameExpr
from mypy.nodes import Statement
from mypy.nodes import SymbolTableNode
from mypy.nodes import TypeAlias
from mypy.nodes import TypeInfo
from mypy.options import Options
from mypy.plugin import ClassDefContext
from mypy.plugin import DynamicClassDefContext
from mypy.plugin import SemanticAnalyzerPluginInterface
from mypy.plugins.common import deserialize_and_fixup_type
from mypy.typeops import map_type_from_supertype
from mypy.types import CallableType
from mypy.types import get_proper_type
from mypy.types import Instance
from mypy.types import NoneType
from mypy.types import Type
from mypy.types import TypeVarType
from mypy.types import UnboundType
from mypy.types import UnionType
_vers = tuple(
[int(x) for x in version.__version__.split(".") if re.match(r"^\d+$", x)]
)
mypy_14 = _vers >= (1, 4)
_TArgType = TypeVar("_TArgType", bound=Union[CallExpr, NameExpr])
class SQLAlchemyAttribute:
def __init__(
self,
name: str,
line: int,
column: int,
typ: Optional[Type],
info: TypeInfo,
) -> None:
self.name = name
self.line = line
self.column = column
self.type = typ
self.info = info
def serialize(self) -> JsonDict:
assert self.type
return {
"name": self.name,
"line": self.line,
"column": self.column,
"type": serialize_type(self.type),
}
def expand_typevar_from_subtype(self, sub_type: TypeInfo) -> None:
"""Expands type vars in the context of a subtype when an attribute is
inherited from a generic super type.
"""
if not isinstance(self.type, TypeVarType):
return
self.type = map_type_from_supertype(self.type, sub_type, self.info)
@classmethod
def deserialize(
cls,
info: TypeInfo,
data: JsonDict,
api: SemanticAnalyzerPluginInterface,
) -> SQLAlchemyAttribute:
data = data.copy()
typ = deserialize_and_fixup_type(data.pop("type"), api)
return cls(typ=typ, info=info, **data)
def name_is_dunder(name: str) -> bool:
return bool(re.match(r"^__.+?__$", name))
def _set_info_metadata(info: TypeInfo, key: str, data: Any) -> None:
info.metadata.setdefault("sqlalchemy", {})[key] = data
def _get_info_metadata(info: TypeInfo, key: str) -> Optional[Any]:
return info.metadata.get("sqlalchemy", {}).get(key, None)
def _get_info_mro_metadata(info: TypeInfo, key: str) -> Optional[Any]:
if info.mro:
for base in info.mro:
metadata = _get_info_metadata(base, key)
if metadata is not None:
return metadata
return None
def establish_as_sqlalchemy(info: TypeInfo) -> None:
info.metadata.setdefault("sqlalchemy", {})
def set_is_base(info: TypeInfo) -> None:
_set_info_metadata(info, "is_base", True)
def get_is_base(info: TypeInfo) -> bool:
is_base = _get_info_metadata(info, "is_base")
return is_base is True
def has_declarative_base(info: TypeInfo) -> bool:
is_base = _get_info_mro_metadata(info, "is_base")
return is_base is True
def set_has_table(info: TypeInfo) -> None:
_set_info_metadata(info, "has_table", True)
def get_has_table(info: TypeInfo) -> bool:
is_base = _get_info_metadata(info, "has_table")
return is_base is True
def get_mapped_attributes(
info: TypeInfo, api: SemanticAnalyzerPluginInterface
) -> Optional[List[SQLAlchemyAttribute]]:
mapped_attributes: Optional[List[JsonDict]] = _get_info_metadata(
info, "mapped_attributes"
)
if mapped_attributes is None:
return None
attributes: List[SQLAlchemyAttribute] = []
for data in mapped_attributes:
attr = SQLAlchemyAttribute.deserialize(info, data, api)
attr.expand_typevar_from_subtype(info)
attributes.append(attr)
return attributes
def format_type(typ_: Type, options: Options) -> str:
if mypy_14:
return _mypy_format_type(typ_, options)
else:
return _mypy_format_type(typ_) # type: ignore
def set_mapped_attributes(
info: TypeInfo, attributes: List[SQLAlchemyAttribute]
) -> None:
_set_info_metadata(
info,
"mapped_attributes",
[attribute.serialize() for attribute in attributes],
)
def fail(api: SemanticAnalyzerPluginInterface, msg: str, ctx: Context) -> None:
msg = "[SQLAlchemy Mypy plugin] %s" % msg
return api.fail(msg, ctx)
def add_global(
ctx: Union[ClassDefContext, DynamicClassDefContext],
module: str,
symbol_name: str,
asname: str,
) -> None:
module_globals = ctx.api.modules[ctx.api.cur_mod_id].names
if asname not in module_globals:
lookup_sym: SymbolTableNode = ctx.api.modules[module].names[
symbol_name
]
module_globals[asname] = lookup_sym
@overload
def get_callexpr_kwarg(
callexpr: CallExpr, name: str, *, expr_types: None = ...
) -> Optional[Union[CallExpr, NameExpr]]: ...
@overload
def get_callexpr_kwarg(
callexpr: CallExpr,
name: str,
*,
expr_types: Tuple[TypingType[_TArgType], ...],
) -> Optional[_TArgType]: ...
def get_callexpr_kwarg(
callexpr: CallExpr,
name: str,
*,
expr_types: Optional[Tuple[TypingType[Any], ...]] = None,
) -> Optional[Any]:
try:
arg_idx = callexpr.arg_names.index(name)
except ValueError:
return None
kwarg = callexpr.args[arg_idx]
if isinstance(
kwarg, expr_types if expr_types is not None else (NameExpr, CallExpr)
):
return kwarg
return None
def flatten_typechecking(stmts: Iterable[Statement]) -> Iterator[Statement]:
for stmt in stmts:
if (
isinstance(stmt, IfStmt)
and isinstance(stmt.expr[0], NameExpr)
and stmt.expr[0].fullname == "typing.TYPE_CHECKING"
):
yield from stmt.body[0].body
else:
yield stmt
def type_for_callee(callee: Expression) -> Optional[Union[Instance, TypeInfo]]:
if isinstance(callee, (MemberExpr, NameExpr)):
if isinstance(callee.node, FuncDef):
if callee.node.type and isinstance(callee.node.type, CallableType):
ret_type = get_proper_type(callee.node.type.ret_type)
if isinstance(ret_type, Instance):
return ret_type
return None
elif isinstance(callee.node, TypeAlias):
target_type = get_proper_type(callee.node.target)
if isinstance(target_type, Instance):
return target_type
elif isinstance(callee.node, TypeInfo):
return callee.node
return None
def unbound_to_instance(
api: SemanticAnalyzerPluginInterface, typ: Type
) -> Type:
"""Take the UnboundType that we seem to get as the ret_type from a FuncDef
and convert it into an Instance/TypeInfo kind of structure that seems
to work as the left-hand type of an AssignmentStatement.
"""
if not isinstance(typ, UnboundType):
return typ
# TODO: figure out a more robust way to check this. The node is some
# kind of _SpecialForm, there's a typing.Optional that's _SpecialForm,
# but I can't figure out how to get them to match up
if typ.name == "Optional":
# convert from "Optional?" to the more familiar
# UnionType[..., NoneType()]
return unbound_to_instance(
api,
UnionType(
[unbound_to_instance(api, typ_arg) for typ_arg in typ.args]
+ [NoneType()]
),
)
node = api.lookup_qualified(typ.name, typ)
if (
node is not None
and isinstance(node, SymbolTableNode)
and isinstance(node.node, TypeInfo)
):
bound_type = node.node
return Instance(
bound_type,
[
(
unbound_to_instance(api, arg)
if isinstance(arg, UnboundType)
else arg
)
for arg in typ.args
],
)
else:
return typ
def info_for_cls(
cls: ClassDef, api: SemanticAnalyzerPluginInterface
) -> Optional[TypeInfo]:
if cls.info is CLASSDEF_NO_INFO:
sym = api.lookup_qualified(cls.name, cls)
if sym is None:
return None
assert sym and isinstance(sym.node, TypeInfo)
return sym.node
return cls.info
def serialize_type(typ: Type) -> Union[str, JsonDict]:
try:
return typ.serialize()
except Exception:
pass
if hasattr(typ, "args"):
typ.args = tuple(
(
a.resolve_string_annotation()
if hasattr(a, "resolve_string_annotation")
else a
)
for a in typ.args
)
elif hasattr(typ, "resolve_string_annotation"):
typ = typ.resolve_string_annotation()
return typ.serialize()