Update mypy plugin to only use public plugin API

Change-Id: Id7f4e4a39e17c1b6ec3c754e2fc5c6ba4b437c38

6 files changed, 40 insertions, 35 deletions

diff --git a/lib/sqlalchemy/ext/mypy/apply.py b/lib/sqlalchemy/ext/mypy/apply.py
index 6442cbc22..0f4bb1fd9 100644
--- a/lib/sqlalchemy/ext/mypy/apply.py
+++ b/lib/sqlalchemy/ext/mypy/apply.py
@@ -83,7 +83,6 @@ def _re_apply_declarative_assignments(
         name: typ for name, typ in cls_metadata.mapped_attr_names
     }
 
-    descriptor = api.lookup("__sa_Mapped", cls)
     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
@@ -96,8 +95,7 @@ def _re_apply_declarative_assignments(
             typ = mapped_attr_lookup[stmt.lvalues[0].name]
             left_node = stmt.lvalues[0].node
 
-            inst = Instance(descriptor.node, [typ])
-            left_node.type = inst
+            left_node.type = api.named_type("__sa_Mapped", [typ])
 
 
 def _apply_type_to_mapped_statement(
@@ -125,16 +123,15 @@ def _apply_type_to_mapped_statement(
             attrname : Mapped[Optional[int]] = <meaningless temp node>
 
     """
-    descriptor = api.lookup("__sa_Mapped", stmt)
     left_node = lvalue.node
 
-    inst = Instance(descriptor.node, [python_type_for_type])
-
     if left_hand_explicit_type is not None:
-        left_node.type = Instance(descriptor.node, [left_hand_explicit_type])
+        left_node.type = api.named_type(
+            "__sa_Mapped", [left_hand_explicit_type]
+        )
     else:
         lvalue.is_inferred_def = False
-        left_node.type = inst
+        left_node.type = api.named_type("__sa_Mapped", [python_type_for_type])
 
     # so to have it skip the right side totally, we can do this:
     # stmt.rvalue = TempNode(AnyType(TypeOfAny.special_form))
@@ -206,7 +203,7 @@ def _apply_placeholder_attr_to_class(
     sym = api.lookup_fully_qualified_or_none(qualified_name)
     if sym:
         assert isinstance(sym.node, TypeInfo)
-        type_ = Instance(sym.node, [])
+        type_: Union[Instance, AnyType] = Instance(sym.node, [])
     else:
         type_ = AnyType(TypeOfAny.special_form)
     var = Var(attrname)
diff --git a/lib/sqlalchemy/ext/mypy/decl_class.py b/lib/sqlalchemy/ext/mypy/decl_class.py
index a0e272f71..40f1f0c0f 100644
--- a/lib/sqlalchemy/ext/mypy/decl_class.py
+++ b/lib/sqlalchemy/ext/mypy/decl_class.py
@@ -131,7 +131,7 @@ def _scan_symbol_table_entry(
                 typeengine_arg = typeengine_arg.type
 
             if isinstance(typeengine_arg, (UnboundType, TypeInfo)):
-                sym = api.lookup(typeengine_arg.name, typeengine_arg)
+                sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
                 if sym is not None:
                     if names._mro_has_id(sym.node.mro, names.TYPEENGINE):
 
@@ -229,7 +229,7 @@ def _scan_declarative_decorator_stmt(
         elif type_id is names.COLUMN and func_type.args:
             typeengine_arg = func_type.args[0]
             if isinstance(typeengine_arg, UnboundType):
-                sym = api.lookup(typeengine_arg.name, typeengine_arg)
+                sym = api.lookup_qualified(typeengine_arg.name, typeengine_arg)
                 if sym is not None and names._mro_has_id(
                     sym.node.mro, names.TYPEENGINE
                 ):
@@ -264,8 +264,6 @@ def _scan_declarative_decorator_stmt(
 
         left_hand_explicit_type = AnyType(TypeOfAny.special_form)
 
-    descriptor = api.lookup("__sa_Mapped", cls)
-
     left_node = NameExpr(stmt.var.name)
     left_node.node = stmt.var
 
@@ -280,7 +278,9 @@ def _scan_declarative_decorator_stmt(
             api, left_hand_explicit_type
         )
 
-    left_node.node.type = Instance(descriptor.node, [left_hand_explicit_type])
+    left_node.node.type = api.named_type(
+        "__sa_Mapped", [left_hand_explicit_type]
+    )
 
     # this will ignore the rvalue entirely
     # rvalue = TempNode(AnyType(TypeOfAny.special_form))
@@ -290,7 +290,7 @@ def _scan_declarative_decorator_stmt(
     # _sa_Mapped._empty_constructor(lambda: <function body>)
     # the function body is maintained so it gets type checked internally
     column_descriptor = nodes.NameExpr("__sa_Mapped")
-    column_descriptor.fullname = "sqlalchemy.orm.Mapped"
+    column_descriptor.fullname = "sqlalchemy.orm.attributes.Mapped"
     mm = nodes.MemberExpr(column_descriptor, "_empty_constructor")
 
     arg = nodes.LambdaExpr(stmt.func.arguments, stmt.func.body)
@@ -367,7 +367,7 @@ def _scan_declarative_assignment_stmt(
             left_hand_explicit_type = stmt.type
 
             if stmt.type.name == "Mapped":
-                mapped_sym = api.lookup("Mapped", cls)
+                mapped_sym = api.lookup_qualified("Mapped", cls)
                 if (
                     mapped_sym is not None
                     and names._type_id_for_named_node(mapped_sym.node)
diff --git a/lib/sqlalchemy/ext/mypy/infer.py b/lib/sqlalchemy/ext/mypy/infer.py
index f0f6be36f..f1bda7865 100644
--- a/lib/sqlalchemy/ext/mypy/infer.py
+++ b/lib/sqlalchemy/ext/mypy/infer.py
@@ -103,9 +103,8 @@ def _infer_type_from_relationship(
     ):
         type_is_a_collection = True
         if python_type_for_type is not None:
-            python_type_for_type = Instance(
-                api.lookup_fully_qualified("builtins.list").node,
-                [python_type_for_type],
+            python_type_for_type = api.named_type(
+                "__builtins__.list", [python_type_for_type]
             )
     elif (
         uselist_arg is None or uselist_arg.fullname == "builtins.True"
@@ -349,9 +348,9 @@ def _infer_type_from_left_and_inferred_right(
         python_type_for_type = python_type_for_type.args[0]
 
     if not is_subtype(left_hand_explicit_type, python_type_for_type):
-        descriptor = api.lookup("__sa_Mapped", node)
-
-        effective_type = Instance(descriptor.node, [orig_python_type_for_type])
+        effective_type = api.named_type(
+            "__sa_Mapped", [orig_python_type_for_type]
+        )
 
         msg = (
             "Left hand assignment '{}: {}' not compatible "
@@ -390,8 +389,7 @@ def _infer_type_from_left_hand_type_only(
         )
         util.fail(api, msg.format(node.name), node)
 
-        descriptor = api.lookup("__sa_Mapped", node)
-        return Instance(descriptor.node, [AnyType(TypeOfAny.special_form)])
+        return api.named_type("__sa_Mapped", [AnyType(TypeOfAny.special_form)])
 
     else:
         # use type from the left hand side
@@ -411,14 +409,21 @@ def _extract_python_type_from_typeengine(
                     return Instance(first_arg.node, [])
             # TODO: support other pep-435 types here
         else:
-            n = api.lookup_fully_qualified("builtins.str")
-            return Instance(n.node, [])
+            return api.named_type("__builtins__.str", [])
 
     assert node.has_base("sqlalchemy.sql.type_api.TypeEngine"), (
         "could not extract Python type from node: %s" % node
     )
+
+    type_engine_sym = api.lookup_fully_qualified_or_none(
+        "sqlalchemy.sql.type_api.TypeEngine"
+    )
+
+    assert type_engine_sym is not None and isinstance(
+        type_engine_sym.node, TypeInfo
+    )
     type_engine = map_instance_to_supertype(
         Instance(node, []),
-        api.modules["sqlalchemy.sql.type_api"].names["TypeEngine"].node,
+        type_engine_sym.node,
     )
     return type_engine.args[-1]
diff --git a/lib/sqlalchemy/ext/mypy/names.py b/lib/sqlalchemy/ext/mypy/names.py
index 11208f3c7..174a8f422 100644
--- a/lib/sqlalchemy/ext/mypy/names.py
+++ b/lib/sqlalchemy/ext/mypy/names.py
@@ -161,7 +161,7 @@ def _type_id_for_unbound_type(
 ) -> int:
     type_id = None
 
-    sym = api.lookup(type_.name, type_)
+    sym = api.lookup_qualified(type_.name, type_)
     if sym is not None:
         if isinstance(sym.node, TypeAlias):
             type_id = _type_id_for_named_node(sym.node.target.type)
diff --git a/lib/sqlalchemy/ext/mypy/plugin.py b/lib/sqlalchemy/ext/mypy/plugin.py
index a0aa5bf04..23585be49 100644
--- a/lib/sqlalchemy/ext/mypy/plugin.py
+++ b/lib/sqlalchemy/ext/mypy/plugin.py
@@ -143,7 +143,9 @@ def _fill_in_decorators(ctx: ClassDefContext) -> None:
         else:
             continue
 
-        sym = ctx.api.lookup(target.expr.name, target, suppress_errors=True)
+        sym = ctx.api.lookup_qualified(
+            target.expr.name, target, suppress_errors=True
+        )
         if sym:
             if sym.node.type and hasattr(sym.node.type, "type"):
                 target.fullname = (
@@ -242,7 +244,7 @@ def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None:
         )
         info.bases = [Instance(cls_arg.node, [])]
     else:
-        obj = ctx.api.builtin_type("builtins.object")
+        obj = ctx.api.named_type("__builtins__.object")
 
         info.bases = [obj]
 
@@ -252,7 +254,7 @@ def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None:
         util.fail(
             ctx.api, "Not able to calculate MRO for declarative base", ctx.call
         )
-        obj = ctx.api.builtin_type("builtins.object")
+        obj = ctx.api.named_type("__builtins__.object")
         info.bases = [obj]
         info.fallback_to_any = True
 
@@ -268,7 +270,7 @@ def _make_declarative_meta(
     declarative_meta_name.fullname = "sqlalchemy.orm.decl_api.DeclarativeMeta"
 
     # installed by _add_globals
-    sym = api.lookup("__sa_DeclarativeMeta", target_cls)
+    sym = api.lookup_qualified("__sa_DeclarativeMeta", target_cls)
 
     declarative_meta_typeinfo = sym.node
     declarative_meta_name.node = declarative_meta_typeinfo
diff --git a/lib/sqlalchemy/ext/mypy/util.py b/lib/sqlalchemy/ext/mypy/util.py
index becce3ebe..1c1e56d2c 100644
--- a/lib/sqlalchemy/ext/mypy/util.py
+++ b/lib/sqlalchemy/ext/mypy/util.py
@@ -11,6 +11,7 @@ from mypy.nodes import JsonDict
 from mypy.nodes import NameExpr
 from mypy.nodes import SymbolTableNode
 from mypy.nodes import TypeInfo
+from mypy.plugin import ClassDefContext
 from mypy.plugin import SemanticAnalyzerPluginInterface
 from mypy.plugins.common import deserialize_and_fixup_type
 from mypy.types import Instance
@@ -68,7 +69,7 @@ def fail(api: SemanticAnalyzerPluginInterface, msg: str, ctx: Context):
 
 
 def add_global(
-    ctx: SemanticAnalyzerPluginInterface,
+    ctx: ClassDefContext,
     module: str,
     symbol_name: str,
     asname: str,
@@ -127,7 +128,7 @@ def _unbound_to_instance(
             ),
         )
 
-    node = api.lookup(typ.name, typ)
+    node = api.lookup_qualified(typ.name, typ)
 
     if node is not None and isinstance(node, SymbolTableNode):
         bound_type = node.node
@@ -147,7 +148,7 @@ def _unbound_to_instance(
 
 def _info_for_cls(cls, api):
     if cls.info is CLASSDEF_NO_INFO:
-        sym = api.lookup(cls.name, cls)
+        sym = api.lookup_qualified(cls.name, cls)
         if sym.node and isinstance(sym.node, TypeInfo):
             info = sym.node
     else: