Fix get_dynamic_class_hook does not work in some scenarios

mypy/semanal.py

@@ -2277,25 +2277,33 @@ def analyze_lvalues(self, s: AssignmentStmt) -> None:
                                 is_final=s.is_final_def)
 
     def apply_dynamic_class_hook(self, s: AssignmentStmt) -> None:
-        if len(s.lvalues) > 1:
-            return
-        lval = s.lvalues[0]
-        if not isinstance(lval, NameExpr) or not isinstance(s.rvalue, CallExpr):
+        if not isinstance(s.rvalue, CallExpr):
             return
-        call = s.rvalue
         fname = None
-        if isinstance(call.callee, RefExpr):
-            fname = call.callee.fullname
-        # check if method call
-        if fname is None and isinstance(call.callee, MemberExpr):
-            callee_expr = call.callee.expr
-            if isinstance(callee_expr, RefExpr) and callee_expr.fullname:
-                method_name = call.callee.name
-                fname = callee_expr.fullname + '.' + method_name
-        if fname:
-            hook = self.plugin.get_dynamic_class_hook(fname)
-            if hook:
-                hook(DynamicClassDefContext(call, lval.name, self))
+        call = s.rvalue
+        while True:
+            if isinstance(call.callee, RefExpr):
+                fname = call.callee.fullname
+            # check if method call
+            if fname is None and isinstance(call.callee, MemberExpr):
+                callee_expr = call.callee.expr
+                if isinstance(callee_expr, RefExpr) and callee_expr.fullname:
+                    method_name = call.callee.name
+                    fname = callee_expr.fullname + '.' + method_name
+                elif isinstance(callee_expr, CallExpr):
+                    # check if chain call
+                    call = callee_expr
+                    continue
+            break
+        if not fname:
+            return
+        hook = self.plugin.get_dynamic_class_hook(fname)
+        if not hook:
+            return
+        for lval in s.lvalues:
+            if not isinstance(lval, NameExpr):
+                continue
+            hook(DynamicClassDefContext(call, lval.name, self))
 
     def unwrap_final(self, s: AssignmentStmt) -> bool:
         """Strip Final[...] if present in an assignment.

test-data/unit/check-custom-plugin.test

@@ -543,26 +543,23 @@ class Instr(Generic[T]): ...
 \[mypy]
 plugins=<ROOT>/test-data/unit/plugins/dyn_class.py
 
-[case testDynamicClassPluginNegatives]
+[case testDynamicClassPluginChainCall]
 # flags: --config-file tmp/mypy.ini
-from mod import declarative_base, Column, Instr, non_declarative_base
+from mod import declarative_base, Column, Instr
 
-Bad1 = non_declarative_base()
-Bad2 = Bad3 = declarative_base()
+Base = declarative_base().with_optional_xxx()
 
-class C1(Bad1): ...  # E: Variable "__main__.Bad1" is not valid as a type \
-                     # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases \
-                     # E: Invalid base class "Bad1"
-class C2(Bad2): ...  # E: Variable "__main__.Bad2" is not valid as a type \
-                     # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases \
-                     # E: Invalid base class "Bad2"
-class C3(Bad3): ...  # E: Variable "__main__.Bad3" is not valid as a type \
-                     # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases \
-                     # E: Invalid base class "Bad3"
+class Model(Base):
+    x: Column[int]
+
+reveal_type(Model().x)  # N: Revealed type is "mod.Instr[builtins.int]"
 [file mod.py]
 from typing import Generic, TypeVar
-def declarative_base(): ...
-def non_declarative_base(): ...
+
+class Base:
+    def with_optional_xxx(self) -> Base: ...
+
+def declarative_base() -> Base: ...
 
 T = TypeVar('T')
 
@@ -573,6 +570,35 @@ class Instr(Generic[T]): ...
 \[mypy]
 plugins=<ROOT>/test-data/unit/plugins/dyn_class.py
 
+[case testDynamicClassPluginChainedAssignment]
+# flags: --config-file tmp/mypy.ini
+from mod import declarative_base
+
+Base1 = Base2 = declarative_base()
+
+class C1(Base1): ...  
+class C2(Base2): ...
+[file mod.py]
+def declarative_base(): ...
+[file mypy.ini]
+\[mypy]
+plugins=<ROOT>/test-data/unit/plugins/dyn_class.py
+
+[case testDynamicClassPluginNegatives]
+# flags: --config-file tmp/mypy.ini
+from mod import non_declarative_base
+
+Bad1 = non_declarative_base()
+
+class C1(Bad1): ...  # E: Variable "__main__.Bad1" is not valid as a type \
+                     # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases \
+                     # E: Invalid base class "Bad1"
+[file mod.py]
+def non_declarative_base(): ...
+[file mypy.ini]
+\[mypy]
+plugins=<ROOT>/test-data/unit/plugins/dyn_class.py
+
 [case testDynamicClassHookFromClassMethod]
 # flags: --config-file tmp/mypy.ini