Merge "Add DeclarativeMeta to globals"

4 files changed, 79 insertions, 31 deletions

diff --git a/lib/sqlalchemy/ext/mypy/decl_class.py b/lib/sqlalchemy/ext/mypy/decl_class.py
index 9fb1fa807..46f3cc30e 100644
--- a/lib/sqlalchemy/ext/mypy/decl_class.py
+++ b/lib/sqlalchemy/ext/mypy/decl_class.py
@@ -229,7 +229,7 @@ def _scan_declarative_decorator_stmt(
 
         left_hand_explicit_type = AnyType(TypeOfAny.special_form)
 
-    descriptor = api.modules["sqlalchemy.orm.attributes"].names["Mapped"]
+    descriptor = api.lookup("__sa_Mapped", cls)
 
     left_node = NameExpr(stmt.var.name)
     left_node.node = stmt.var
@@ -254,8 +254,7 @@ def _scan_declarative_decorator_stmt(
     # <attr> : Mapped[<typ>] =
     # _sa_Mapped._empty_constructor(lambda: <function body>)
     # the function body is maintained so it gets type checked internally
-    api.add_symbol_table_node("_sa_Mapped", descriptor)
-    column_descriptor = nodes.NameExpr("_sa_Mapped")
+    column_descriptor = nodes.NameExpr("__sa_Mapped")
     column_descriptor.fullname = "sqlalchemy.orm.Mapped"
     mm = nodes.MemberExpr(column_descriptor, "_empty_constructor")
 
@@ -753,7 +752,7 @@ def _infer_type_from_left_and_inferred_right(
 
     """
     if not is_subtype(left_hand_explicit_type, python_type_for_type):
-        descriptor = api.modules["sqlalchemy.orm.attributes"].names["Mapped"]
+        descriptor = api.lookup("__sa_Mapped", node)
 
         effective_type = Instance(descriptor.node, [python_type_for_type])
 
@@ -794,8 +793,7 @@ def _infer_type_from_left_hand_type_only(
         )
         util.fail(api, msg.format(node.name), node)
 
-        descriptor = api.modules["sqlalchemy.orm.attributes"].names["Mapped"]
-
+        descriptor = api.lookup("__sa_Mapped", node)
         return Instance(descriptor.node, [AnyType(TypeOfAny.special_form)])
 
     else:
@@ -816,8 +814,7 @@ def _re_apply_declarative_assignments(
         name: typ for name, typ in cls_metadata.mapped_attr_names
     }
 
-    descriptor = api.modules["sqlalchemy.orm.attributes"].names["Mapped"]
-
+    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
@@ -859,8 +856,7 @@ def _apply_type_to_mapped_statement(
             attrname : Mapped[Optional[int]] = <meaningless temp node>
 
     """
-    descriptor = api.modules["sqlalchemy.orm.attributes"].names["Mapped"]
-
+    descriptor = api.lookup("__sa_Mapped", stmt)
     left_node = lvalue.node
 
     inst = Instance(descriptor.node, [python_type_for_type])
diff --git a/lib/sqlalchemy/ext/mypy/plugin.py b/lib/sqlalchemy/ext/mypy/plugin.py
index c8fbcd6a2..9ca1cb2da 100644
--- a/lib/sqlalchemy/ext/mypy/plugin.py
+++ b/lib/sqlalchemy/ext/mypy/plugin.py
@@ -126,6 +126,7 @@ def _fill_in_decorators(ctx: ClassDefContext) -> None:
         # 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)
@@ -165,15 +166,34 @@ def _fill_in_decorators(ctx: ClassDefContext) -> None:
                 )
 
 
+def _add_globals(ctx: ClassDefContext):
+    """Add __sa_DeclarativeMeta and __sa_Mapped symbol to the global space
+    for all class defs
+
+    """
+
+    util.add_global(
+        ctx,
+        "sqlalchemy.orm.decl_api",
+        "DeclarativeMeta",
+        "__sa_DeclarativeMeta",
+    )
+
+    util.add_global(ctx, "sqlalchemy.orm.attributes", "Mapped", "__sa_Mapped")
+
+
 def _cls_metadata_hook(ctx: ClassDefContext) -> None:
+    _add_globals(ctx)
     decl_class._scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
 
 
 def _base_cls_hook(ctx: ClassDefContext) -> None:
+    _add_globals(ctx)
     decl_class._scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
 
 
 def _cls_decorator_hook(ctx: ClassDefContext) -> None:
+    _add_globals(ctx)
     assert isinstance(ctx.reason, nodes.MemberExpr)
     expr = ctx.reason.expr
     assert names._type_id_for_named_node(expr.node.type.type) is names.REGISTRY
@@ -181,28 +201,8 @@ def _cls_decorator_hook(ctx: ClassDefContext) -> None:
     decl_class._scan_declarative_assignments_and_apply_types(ctx.cls, ctx.api)
 
 
-def _make_declarative_meta(
-    api: SemanticAnalyzerPluginInterface, target_cls: ClassDef
-):
-    declarative_meta_sym: SymbolTableNode = api.modules[
-        "sqlalchemy.orm.decl_api"
-    ].names["DeclarativeMeta"]
-    declarative_meta_typeinfo: TypeInfo = declarative_meta_sym.node
-
-    declarative_meta_name: NameExpr = NameExpr("DeclarativeMeta")
-    declarative_meta_name.kind = GDEF
-    declarative_meta_name.fullname = "sqlalchemy.orm.decl_api.DeclarativeMeta"
-    declarative_meta_name.node = declarative_meta_typeinfo
-
-    target_cls.metaclass = declarative_meta_name
-
-    declarative_meta_instance = Instance(declarative_meta_typeinfo, [])
-
-    info = target_cls.info
-    info.declared_metaclass = info.metaclass_type = declarative_meta_instance
-
-
 def _base_cls_decorator_hook(ctx: ClassDefContext) -> None:
+    _add_globals(ctx)
 
     cls = ctx.cls
 
@@ -217,6 +217,8 @@ 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)
 
@@ -246,3 +248,25 @@ def _dynamic_class_hook(ctx: DynamicClassDefContext) -> None:
         info.fallback_to_any = True
 
     ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, info))
+
+
+def _make_declarative_meta(
+    api: SemanticAnalyzerPluginInterface, target_cls: ClassDef
+):
+
+    declarative_meta_name: NameExpr = NameExpr("__sa_DeclarativeMeta")
+    declarative_meta_name.kind = GDEF
+    declarative_meta_name.fullname = "sqlalchemy.orm.decl_api.DeclarativeMeta"
+
+    # installed by _add_globals
+    sym = api.lookup("__sa_DeclarativeMeta", target_cls)
+
+    declarative_meta_typeinfo = sym.node
+    declarative_meta_name.node = declarative_meta_typeinfo
+
+    target_cls.metaclass = declarative_meta_name
+
+    declarative_meta_instance = Instance(declarative_meta_typeinfo, [])
+
+    info = target_cls.info
+    info.declared_metaclass = info.metaclass_type = declarative_meta_instance
diff --git a/lib/sqlalchemy/ext/mypy/util.py b/lib/sqlalchemy/ext/mypy/util.py
index e7178a885..7079f3cd7 100644
--- a/lib/sqlalchemy/ext/mypy/util.py
+++ b/lib/sqlalchemy/ext/mypy/util.py
@@ -18,6 +18,22 @@ def fail(api: SemanticAnalyzerPluginInterface, msg: str, ctx: Context):
     return api.fail(msg, ctx)
 
 
+def add_global(
+    ctx: SemanticAnalyzerPluginInterface,
+    module: str,
+    symbol_name: str,
+    asname: str,
+):
+    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
+
+
 def _get_callexpr_kwarg(callexpr: CallExpr, name: str) -> Optional[NameExpr]:
     try:
         arg_idx = callexpr.arg_names.index(name)
diff --git a/lib/sqlalchemy/testing/requirements.py b/lib/sqlalchemy/testing/requirements.py
index 46844803b..b6381dd57 100644
--- a/lib/sqlalchemy/testing/requirements.py
+++ b/lib/sqlalchemy/testing/requirements.py
@@ -1236,6 +1236,18 @@ class SuiteRequirements(Requirements):
         )
 
     @property
+    def patch_library(self):
+        def check_lib():
+            try:
+                __import__("patch")
+            except ImportError:
+                return False
+            else:
+                return True
+
+        return exclusions.only_if(check_lib, "patch library needed")
+
+    @property
     def non_broken_pickle(self):
         from sqlalchemy.util import pickle