SQLAlchemy 2.0.47__cp313-cp313t-win32.whl

sqlalchemy/ext/mypy/decl_class.py

@@ -0,0 +1,515 @@
+# ext/mypy/decl_class.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
+
+from typing import List
+from typing import Optional
+from typing import Union
+
+from mypy.nodes import AssignmentStmt
+from mypy.nodes import CallExpr
+from mypy.nodes import ClassDef
+from mypy.nodes import Decorator
+from mypy.nodes import LambdaExpr
+from mypy.nodes import ListExpr
+from mypy.nodes import MemberExpr
+from mypy.nodes import NameExpr
+from mypy.nodes import PlaceholderNode
+from mypy.nodes import RefExpr
+from mypy.nodes import StrExpr
+from mypy.nodes import SymbolNode
+from mypy.nodes import SymbolTableNode
+from mypy.nodes import TempNode
+from mypy.nodes import TypeInfo
+from mypy.nodes import Var
+from mypy.plugin import SemanticAnalyzerPluginInterface
+from mypy.types import AnyType
+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 ProperType
+from mypy.types import Type
+from mypy.types import TypeOfAny
+from mypy.types import UnboundType
+from mypy.types import UnionType
+
+from . import apply
+from . import infer
+from . import names
+from . import util
+
+
+def scan_declarative_assignments_and_apply_types(
+    cls: ClassDef,
+    api: SemanticAnalyzerPluginInterface,
+    is_mixin_scan: bool = False,
+) -> Optional[List[util.SQLAlchemyAttribute]]:
+    info = util.info_for_cls(cls, api)
+
+    if info is None:
+        # this can occur during cached passes
+        return None
+    elif cls.fullname.startswith("builtins"):
+        return None
+
+    mapped_attributes: Optional[List[util.SQLAlchemyAttribute]] = (
+        util.get_mapped_attributes(info, api)
+    )
+
+    # used by assign.add_additional_orm_attributes among others
+    util.establish_as_sqlalchemy(info)
+
+    if mapped_attributes is not None:
+        # ensure that a class that's mapped is always picked up by
+        # its mapped() decorator or declarative metaclass before
+        # it would be detected as an unmapped mixin class
+
+        if not is_mixin_scan:
+            # mypy can call us more than once.  it then *may* have reset the
+            # left hand side of everything, but not the right that we removed,
+            # removing our ability to re-scan.   but we have the types
+            # here, so lets re-apply them, or if we have an UnboundType,
+            # we can re-scan
+
+            apply.re_apply_declarative_assignments(cls, api, mapped_attributes)
+
+        return mapped_attributes
+
+    mapped_attributes = []
+
+    if not cls.defs.body:
+        # when we get a mixin class from another file, the body is
+        # empty (!) but the names are in the symbol table.  so use that.
+
+        for sym_name, sym in info.names.items():
+            _scan_symbol_table_entry(
+                cls, api, sym_name, sym, mapped_attributes
+            )
+    else:
+        for stmt in util.flatten_typechecking(cls.defs.body):
+            if isinstance(stmt, AssignmentStmt):
+                _scan_declarative_assignment_stmt(
+                    cls, api, stmt, mapped_attributes
+                )
+            elif isinstance(stmt, Decorator):
+                _scan_declarative_decorator_stmt(
+                    cls, api, stmt, mapped_attributes
+                )
+    _scan_for_mapped_bases(cls, api)
+
+    if not is_mixin_scan:
+        apply.add_additional_orm_attributes(cls, api, mapped_attributes)
+
+    util.set_mapped_attributes(info, mapped_attributes)
+
+    return mapped_attributes
+
+
+def _scan_symbol_table_entry(
+    cls: ClassDef,
+    api: SemanticAnalyzerPluginInterface,
+    name: str,
+    value: SymbolTableNode,
+    attributes: List[util.SQLAlchemyAttribute],
+) -> None:
+    """Extract mapping information from a SymbolTableNode that's in the
+    type.names dictionary.
+
+    """
+    value_type = get_proper_type(value.type)
+    if not isinstance(value_type, Instance):
+        return
+
+    left_hand_explicit_type = None
+    type_id = names.type_id_for_named_node(value_type.type)
+    # type_id = names._type_id_for_unbound_type(value.type.type, cls, api)
+
+    err = False
+
+    # TODO: this is nearly the same logic as that of
+    # _scan_declarative_decorator_stmt, likely can be merged
+    if type_id in {
+        names.MAPPED,
+        names.RELATIONSHIP,
+        names.COMPOSITE_PROPERTY,
+        names.MAPPER_PROPERTY,
+        names.SYNONYM_PROPERTY,
+        names.COLUMN_PROPERTY,
+    }:
+        if value_type.args:
+            left_hand_explicit_type = get_proper_type(value_type.args[0])
+        else:
+            err = True
+    elif type_id is names.COLUMN:
+        if not value_type.args:
+            err = True
+        else:
+            typeengine_arg: Union[ProperType, TypeInfo] = get_proper_type(
+                value_type.args[0]
+            )
+            if isinstance(typeengine_arg, Instance):
+                typeengine_arg = typeengine_arg.type
+
+            if isinstance(typeengine_arg, (UnboundType, TypeInfo)):
+                sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
+                if sym is not None and isinstance(sym.node, TypeInfo):
+                    if names.has_base_type_id(sym.node, names.TYPEENGINE):
+                        left_hand_explicit_type = UnionType(
+                            [
+                                infer.extract_python_type_from_typeengine(
+                                    api, sym.node, []
+                                ),
+                                NoneType(),
+                            ]
+                        )
+                    else:
+                        util.fail(
+                            api,
+                            "Column type should be a TypeEngine "
+                            "subclass not '{}'".format(sym.node.fullname),
+                            value_type,
+                        )
+
+    if err:
+        msg = (
+            "Can't infer type from attribute {} on class {}. "
+            "please specify a return type from this function that is "
+            "one of: Mapped[<python type>], relationship[<target class>], "
+            "Column[<TypeEngine>], MapperProperty[<python type>]"
+        )
+        util.fail(api, msg.format(name, cls.name), cls)
+
+        left_hand_explicit_type = AnyType(TypeOfAny.special_form)
+
+    if left_hand_explicit_type is not None:
+        assert value.node is not None
+        attributes.append(
+            util.SQLAlchemyAttribute(
+                name=name,
+                line=value.node.line,
+                column=value.node.column,
+                typ=left_hand_explicit_type,
+                info=cls.info,
+            )
+        )
+
+
+def _scan_declarative_decorator_stmt(
+    cls: ClassDef,
+    api: SemanticAnalyzerPluginInterface,
+    stmt: Decorator,
+    attributes: List[util.SQLAlchemyAttribute],
+) -> None:
+    """Extract mapping information from a @declared_attr in a declarative
+    class.
+
+    E.g.::
+
+        @reg.mapped
+        class MyClass:
+            # ...
+
+            @declared_attr
+            def updated_at(cls) -> Column[DateTime]:
+                return Column(DateTime)
+
+    Will resolve in mypy as::
+
+        @reg.mapped
+        class MyClass:
+            # ...
+
+            updated_at: Mapped[Optional[datetime.datetime]]
+
+    """
+    for dec in stmt.decorators:
+        if (
+            isinstance(dec, (NameExpr, MemberExpr, SymbolNode))
+            and names.type_id_for_named_node(dec) is names.DECLARED_ATTR
+        ):
+            break
+    else:
+        return
+
+    dec_index = cls.defs.body.index(stmt)
+
+    left_hand_explicit_type: Optional[ProperType] = None
+
+    if util.name_is_dunder(stmt.name):
+        # for dunder names like __table_args__, __tablename__,
+        # __mapper_args__ etc., rewrite these as simple assignment
+        # statements; otherwise mypy doesn't like if the decorated
+        # function has an annotation like ``cls: Type[Foo]`` because
+        # it isn't @classmethod
+        any_ = AnyType(TypeOfAny.special_form)
+        left_node = NameExpr(stmt.var.name)
+        left_node.node = stmt.var
+        new_stmt = AssignmentStmt([left_node], TempNode(any_))
+        new_stmt.type = left_node.node.type
+        cls.defs.body[dec_index] = new_stmt
+        return
+    elif isinstance(stmt.func.type, CallableType):
+        func_type = stmt.func.type.ret_type
+        if isinstance(func_type, UnboundType):
+            type_id = names.type_id_for_unbound_type(func_type, cls, api)
+        else:
+            # this does not seem to occur unless the type argument is
+            # incorrect
+            return
+
+        if (
+            type_id
+            in {
+                names.MAPPED,
+                names.RELATIONSHIP,
+                names.COMPOSITE_PROPERTY,
+                names.MAPPER_PROPERTY,
+                names.SYNONYM_PROPERTY,
+                names.COLUMN_PROPERTY,
+            }
+            and func_type.args
+        ):
+            left_hand_explicit_type = get_proper_type(func_type.args[0])
+        elif type_id is names.COLUMN and func_type.args:
+            typeengine_arg = func_type.args[0]
+            if isinstance(typeengine_arg, UnboundType):
+                sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
+                if sym is not None and isinstance(sym.node, TypeInfo):
+                    if names.has_base_type_id(sym.node, names.TYPEENGINE):
+                        left_hand_explicit_type = UnionType(
+                            [
+                                infer.extract_python_type_from_typeengine(
+                                    api, sym.node, []
+                                ),
+                                NoneType(),
+                            ]
+                        )
+                    else:
+                        util.fail(
+                            api,
+                            "Column type should be a TypeEngine "
+                            "subclass not '{}'".format(sym.node.fullname),
+                            func_type,
+                        )
+
+    if left_hand_explicit_type is None:
+        # no type on the decorated function.  our option here is to
+        # dig into the function body and get the return type, but they
+        # should just have an annotation.
+        msg = (
+            "Can't infer type from @declared_attr on function '{}';  "
+            "please specify a return type from this function that is "
+            "one of: Mapped[<python type>], relationship[<target class>], "
+            "Column[<TypeEngine>], MapperProperty[<python type>]"
+        )
+        util.fail(api, msg.format(stmt.var.name), stmt)
+
+        left_hand_explicit_type = AnyType(TypeOfAny.special_form)
+
+    left_node = NameExpr(stmt.var.name)
+    left_node.node = stmt.var
+
+    # totally feeling around in the dark here as I don't totally understand
+    # the significance of UnboundType.  It seems to be something that is
+    # not going to do what's expected when it is applied as the type of
+    # an AssignmentStatement.  So do a feeling-around-in-the-dark version
+    # of converting it to the regular Instance/TypeInfo/UnionType structures
+    # we see everywhere else.
+    if isinstance(left_hand_explicit_type, UnboundType):
+        left_hand_explicit_type = get_proper_type(
+            util.unbound_to_instance(api, left_hand_explicit_type)
+        )
+
+    left_node.node.type = api.named_type(
+        names.NAMED_TYPE_SQLA_MAPPED, [left_hand_explicit_type]
+    )
+
+    # this will ignore the rvalue entirely
+    # rvalue = TempNode(AnyType(TypeOfAny.special_form))
+
+    # rewrite the node as:
+    # <attr> : Mapped[<typ>] =
+    # _sa_Mapped._empty_constructor(lambda: <function body>)
+    # the function body is maintained so it gets type checked internally
+    rvalue = names.expr_to_mapped_constructor(
+        LambdaExpr(stmt.func.arguments, stmt.func.body)
+    )
+
+    new_stmt = AssignmentStmt([left_node], rvalue)
+    new_stmt.type = left_node.node.type
+
+    attributes.append(
+        util.SQLAlchemyAttribute(
+            name=left_node.name,
+            line=stmt.line,
+            column=stmt.column,
+            typ=left_hand_explicit_type,
+            info=cls.info,
+        )
+    )
+    cls.defs.body[dec_index] = new_stmt
+
+
+def _scan_declarative_assignment_stmt(
+    cls: ClassDef,
+    api: SemanticAnalyzerPluginInterface,
+    stmt: AssignmentStmt,
+    attributes: List[util.SQLAlchemyAttribute],
+) -> None:
+    """Extract mapping information from an assignment statement in a
+    declarative class.
+
+    """
+    lvalue = stmt.lvalues[0]
+    if not isinstance(lvalue, NameExpr):
+        return
+
+    sym = cls.info.names.get(lvalue.name)
+
+    # this establishes that semantic analysis has taken place, which
+    # means the nodes are populated and we are called from an appropriate
+    # hook.
+    assert sym is not None
+    node = sym.node
+
+    if isinstance(node, PlaceholderNode):
+        return
+
+    assert node is lvalue.node
+    assert isinstance(node, Var)
+
+    if node.name == "__abstract__":
+        if api.parse_bool(stmt.rvalue) is True:
+            util.set_is_base(cls.info)
+        return
+    elif node.name == "__tablename__":
+        util.set_has_table(cls.info)
+    elif node.name.startswith("__"):
+        return
+    elif node.name == "_mypy_mapped_attrs":
+        if not isinstance(stmt.rvalue, ListExpr):
+            util.fail(api, "_mypy_mapped_attrs is expected to be a list", stmt)
+        else:
+            for item in stmt.rvalue.items:
+                if isinstance(item, (NameExpr, StrExpr)):
+                    apply.apply_mypy_mapped_attr(cls, api, item, attributes)
+
+    left_hand_mapped_type: Optional[Type] = None
+    left_hand_explicit_type: Optional[ProperType] = None
+
+    if node.is_inferred or node.type is None:
+        if isinstance(stmt.type, UnboundType):
+            # look for an explicit Mapped[] type annotation on the left
+            # side with nothing on the right
+
+            # print(stmt.type)
+            # Mapped?[Optional?[A?]]
+
+            left_hand_explicit_type = stmt.type
+
+            if stmt.type.name == "Mapped":
+                mapped_sym = api.lookup_qualified("Mapped", cls)
+                if (
+                    mapped_sym is not None
+                    and mapped_sym.node is not None
+                    and names.type_id_for_named_node(mapped_sym.node)
+                    is names.MAPPED
+                ):
+                    left_hand_explicit_type = get_proper_type(
+                        stmt.type.args[0]
+                    )
+                    left_hand_mapped_type = stmt.type
+
+            # TODO: do we need to convert from unbound for this case?
+            # left_hand_explicit_type = util._unbound_to_instance(
+            #     api, left_hand_explicit_type
+            # )
+    else:
+        node_type = get_proper_type(node.type)
+        if (
+            isinstance(node_type, Instance)
+            and names.type_id_for_named_node(node_type.type) is names.MAPPED
+        ):
+            # print(node.type)
+            # sqlalchemy.orm.attributes.Mapped[<python type>]
+            left_hand_explicit_type = get_proper_type(node_type.args[0])
+            left_hand_mapped_type = node_type
+        else:
+            # print(node.type)
+            # <python type>
+            left_hand_explicit_type = node_type
+            left_hand_mapped_type = None
+
+    if isinstance(stmt.rvalue, TempNode) and left_hand_mapped_type is not None:
+        # annotation without assignment and Mapped is present
+        # as type annotation
+        # equivalent to using _infer_type_from_left_hand_type_only.
+
+        python_type_for_type = left_hand_explicit_type
+    elif isinstance(stmt.rvalue, CallExpr) and isinstance(
+        stmt.rvalue.callee, RefExpr
+    ):
+        python_type_for_type = infer.infer_type_from_right_hand_nameexpr(
+            api, stmt, node, left_hand_explicit_type, stmt.rvalue.callee
+        )
+
+        if python_type_for_type is None:
+            return
+
+    else:
+        return
+
+    assert python_type_for_type is not None
+
+    attributes.append(
+        util.SQLAlchemyAttribute(
+            name=node.name,
+            line=stmt.line,
+            column=stmt.column,
+            typ=python_type_for_type,
+            info=cls.info,
+        )
+    )
+
+    apply.apply_type_to_mapped_statement(
+        api,
+        stmt,
+        lvalue,
+        left_hand_explicit_type,
+        python_type_for_type,
+    )
+
+
+def _scan_for_mapped_bases(
+    cls: ClassDef,
+    api: SemanticAnalyzerPluginInterface,
+) -> None:
+    """Given a class, iterate through its superclass hierarchy to find
+    all other classes that are considered as ORM-significant.
+
+    Locates non-mapped mixins and scans them for mapped attributes to be
+    applied to subclasses.
+
+    """
+
+    info = util.info_for_cls(cls, api)
+
+    if info is None:
+        return
+
+    for base_info in info.mro[1:-1]:
+        if base_info.fullname.startswith("builtins"):
+            continue
+
+        # scan each base for mapped attributes.  if they are not already
+        # scanned (but have all their type info), that means they are unmapped
+        # mixins
+        scan_declarative_assignments_and_apply_types(
+            base_info.defn, api, is_mixin_scan=True
+        )