first add files

1 files changed, 299 insertions, 0 deletions

diff --git a/lib/sqlalchemy/ext/mypy/apply.py b/lib/sqlalchemy/ext/mypy/apply.py
new file mode 100644
index 0000000..99be194
--- /dev/null
+++ b/lib/sqlalchemy/ext/mypy/apply.py
@@ -0,0 +1,299 @@
+# ext/mypy/apply.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
+
+from typing import List
+from typing import Optional
+from typing import Union
+
+from mypy.nodes import ARG_NAMED_OPT
+from mypy.nodes import Argument
+from mypy.nodes import AssignmentStmt
+from mypy.nodes import CallExpr
+from mypy.nodes import ClassDef
+from mypy.nodes import MDEF
+from mypy.nodes import MemberExpr
+from mypy.nodes import NameExpr
+from mypy.nodes import RefExpr
+from mypy.nodes import StrExpr
+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.plugins.common import add_method_to_class
+from mypy.types import AnyType
+from mypy.types import get_proper_type
+from mypy.types import Instance
+from mypy.types import NoneTyp
+from mypy.types import ProperType
+from mypy.types import TypeOfAny
+from mypy.types import UnboundType
+from mypy.types import UnionType
+
+from . import infer
+from . import util
+from .names import NAMED_TYPE_SQLA_MAPPED
+
+
+def apply_mypy_mapped_attr(
+    cls: ClassDef,
+    api: SemanticAnalyzerPluginInterface,
+    item: Union[NameExpr, StrExpr],
+    attributes: List[util.SQLAlchemyAttribute],
+) -> None:
+    if isinstance(item, NameExpr):
+        name = item.name
+    elif isinstance(item, StrExpr):
+        name = item.value
+    else:
+        return None
+
+    for stmt in cls.defs.body:
+        if (
+            isinstance(stmt, AssignmentStmt)
+            and isinstance(stmt.lvalues[0], NameExpr)
+            and stmt.lvalues[0].name == name
+        ):
+            break
+    else:
+        util.fail(api, "Can't find mapped attribute {}".format(name), cls)
+        return None
+
+    if stmt.type is None:
+        util.fail(
+            api,
+            "Statement linked from _mypy_mapped_attrs has no "
+            "typing information",
+            stmt,
+        )
+        return None
+
+    left_hand_explicit_type = get_proper_type(stmt.type)
+    assert isinstance(
+        left_hand_explicit_type, (Instance, UnionType, UnboundType)
+    )
+
+    attributes.append(
+        util.SQLAlchemyAttribute(
+            name=name,
+            line=item.line,
+            column=item.column,
+            typ=left_hand_explicit_type,
+            info=cls.info,
+        )
+    )
+
+    apply_type_to_mapped_statement(
+        api, stmt, stmt.lvalues[0], left_hand_explicit_type, None
+    )
+
+
+def re_apply_declarative_assignments(
+    cls: ClassDef,
+    api: SemanticAnalyzerPluginInterface,
+    attributes: List[util.SQLAlchemyAttribute],
+) -> None:
+    """For multiple class passes, re-apply our left-hand side types as mypy
+    seems to reset them in place.
+
+    """
+    mapped_attr_lookup = {attr.name: attr for attr in attributes}
+    update_cls_metadata = False
+
+    for stmt in cls.defs.body:
+        # for a re-apply, all of our statements are AssignmentStmt;
+        # @declared_attr calls will have been converted and this
+        # currently seems to be preserved by mypy (but who knows if this
+        # will change).
+        if (
+            isinstance(stmt, AssignmentStmt)
+            and isinstance(stmt.lvalues[0], NameExpr)
+            and stmt.lvalues[0].name in mapped_attr_lookup
+            and isinstance(stmt.lvalues[0].node, Var)
+        ):
+
+            left_node = stmt.lvalues[0].node
+            python_type_for_type = mapped_attr_lookup[
+                stmt.lvalues[0].name
+            ].type
+
+            left_node_proper_type = get_proper_type(left_node.type)
+
+            # if we have scanned an UnboundType and now there's a more
+            # specific type than UnboundType, call the re-scan so we
+            # can get that set up correctly
+            if (
+                isinstance(python_type_for_type, UnboundType)
+                and not isinstance(left_node_proper_type, UnboundType)
+                and (
+                    isinstance(stmt.rvalue, CallExpr)
+                    and isinstance(stmt.rvalue.callee, MemberExpr)
+                    and isinstance(stmt.rvalue.callee.expr, NameExpr)
+                    and stmt.rvalue.callee.expr.node is not None
+                    and stmt.rvalue.callee.expr.node.fullname
+                    == NAMED_TYPE_SQLA_MAPPED
+                    and stmt.rvalue.callee.name == "_empty_constructor"
+                    and isinstance(stmt.rvalue.args[0], CallExpr)
+                    and isinstance(stmt.rvalue.args[0].callee, RefExpr)
+                )
+            ):
+
+                python_type_for_type = (
+                    infer.infer_type_from_right_hand_nameexpr(
+                        api,
+                        stmt,
+                        left_node,
+                        left_node_proper_type,
+                        stmt.rvalue.args[0].callee,
+                    )
+                )
+
+                if python_type_for_type is None or isinstance(
+                    python_type_for_type, UnboundType
+                ):
+                    continue
+
+                # update the SQLAlchemyAttribute with the better information
+                mapped_attr_lookup[
+                    stmt.lvalues[0].name
+                ].type = python_type_for_type
+
+                update_cls_metadata = True
+
+            if python_type_for_type is not None:
+                left_node.type = api.named_type(
+                    NAMED_TYPE_SQLA_MAPPED, [python_type_for_type]
+                )
+
+    if update_cls_metadata:
+        util.set_mapped_attributes(cls.info, attributes)
+
+
+def apply_type_to_mapped_statement(
+    api: SemanticAnalyzerPluginInterface,
+    stmt: AssignmentStmt,
+    lvalue: NameExpr,
+    left_hand_explicit_type: Optional[ProperType],
+    python_type_for_type: Optional[ProperType],
+) -> None:
+    """Apply the Mapped[<type>] annotation and right hand object to a
+    declarative assignment statement.
+
+    This converts a Python declarative class statement such as::
+
+        class User(Base):
+            # ...
+
+            attrname = Column(Integer)
+
+    To one that describes the final Python behavior to Mypy::
+
+        class User(Base):
+            # ...
+
+            attrname : Mapped[Optional[int]] = <meaningless temp node>
+
+    """
+    left_node = lvalue.node
+    assert isinstance(left_node, Var)
+
+    if left_hand_explicit_type is not None:
+        left_node.type = api.named_type(
+            NAMED_TYPE_SQLA_MAPPED, [left_hand_explicit_type]
+        )
+    else:
+        lvalue.is_inferred_def = False
+        left_node.type = api.named_type(
+            NAMED_TYPE_SQLA_MAPPED,
+            [] if python_type_for_type is None else [python_type_for_type],
+        )
+
+    # so to have it skip the right side totally, we can do this:
+    # stmt.rvalue = TempNode(AnyType(TypeOfAny.special_form))
+
+    # however, if we instead manufacture a new node that uses the old
+    # one, then we can still get type checking for the call itself,
+    # e.g. the Column, relationship() call, etc.
+
+    # rewrite the node as:
+    # <attr> : Mapped[<typ>] =
+    # _sa_Mapped._empty_constructor(<original CallExpr from rvalue>)
+    # the original right-hand side is maintained so it gets type checked
+    # internally
+    stmt.rvalue = util.expr_to_mapped_constructor(stmt.rvalue)
+
+
+def add_additional_orm_attributes(
+    cls: ClassDef,
+    api: SemanticAnalyzerPluginInterface,
+    attributes: List[util.SQLAlchemyAttribute],
+) -> None:
+    """Apply __init__, __table__ and other attributes to the mapped class."""
+
+    info = util.info_for_cls(cls, api)
+
+    if info is None:
+        return
+
+    is_base = util.get_is_base(info)
+
+    if "__init__" not in info.names and not is_base:
+        mapped_attr_names = {attr.name: attr.type for attr in attributes}
+
+        for base in info.mro[1:-1]:
+            if "sqlalchemy" not in info.metadata:
+                continue
+
+            base_cls_attributes = util.get_mapped_attributes(base, api)
+            if base_cls_attributes is None:
+                continue
+
+            for attr in base_cls_attributes:
+                mapped_attr_names.setdefault(attr.name, attr.type)
+
+        arguments = []
+        for name, typ in mapped_attr_names.items():
+            if typ is None:
+                typ = AnyType(TypeOfAny.special_form)
+            arguments.append(
+                Argument(
+                    variable=Var(name, typ),
+                    type_annotation=typ,
+                    initializer=TempNode(typ),
+                    kind=ARG_NAMED_OPT,
+                )
+            )
+
+        add_method_to_class(api, cls, "__init__", arguments, NoneTyp())
+
+    if "__table__" not in info.names and util.get_has_table(info):
+        _apply_placeholder_attr_to_class(
+            api, cls, "sqlalchemy.sql.schema.Table", "__table__"
+        )
+    if not is_base:
+        _apply_placeholder_attr_to_class(
+            api, cls, "sqlalchemy.orm.mapper.Mapper", "__mapper__"
+        )
+
+
+def _apply_placeholder_attr_to_class(
+    api: SemanticAnalyzerPluginInterface,
+    cls: ClassDef,
+    qualified_name: str,
+    attrname: str,
+) -> None:
+    sym = api.lookup_fully_qualified_or_none(qualified_name)
+    if sym:
+        assert isinstance(sym.node, TypeInfo)
+        type_: ProperType = Instance(sym.node, [])
+    else:
+        type_ = AnyType(TypeOfAny.special_form)
+    var = Var(attrname)
+    var._fullname = cls.fullname + "." + attrname
+    var.info = cls.info
+    var.type = type_
+    cls.info.names[attrname] = SymbolTableNode(MDEF, var)