Handle nested definitions

Author: ilevkivskyi

mypy/build.py

@@ -103,12 +103,12 @@
     ImportBase,
     ImportFrom,
     MypyFile,
-    SymbolNode,
     SymbolTable,
 )
 from mypy.partially_defined import PossiblyUndefinedVariableVisitor
 from mypy.semanal import SemanticAnalyzer
 from mypy.semanal_pass1 import SemanticAnalyzerPreAnalysis
+from mypy.traverser import find_definitions
 from mypy.util import (
     DecodeError,
     decode_python_encoding,
@@ -1085,22 +1085,20 @@ def resolve_location(self, graph: dict[str, State], fullname: str) -> Context |
                 source = decode_python_encoding(state.manager.fscache.read(path))
             tree = parse(source, state.path, state.id, state.manager.errors, state.options)
             self.extra_trees[state.id] = tree
-        defs = tree.defs
+        statements = tree.defs
         while prefix:
             part = prefix.pop(0)
-            for defn in defs:
-                if not isinstance(defn, ClassDef):
+            for statement in statements:
+                defs = find_definitions(statement, part)
+                if not defs or not isinstance((defn := defs[0]), ClassDef):
                     continue
-                if defn.name == part:
-                    defs = defn.defs.body
-                    break
+                statements = defn.defs.body
+                break
             else:
                 return None
-        for defn in defs:
-            # TODO: support more kinds of locations (like assignment statements).
-            # the latter will be helpful for type old-style aliases.
-            if isinstance(defn, SymbolNode) and defn.name == name:
-                return defn
+        for statement in statements:
+            if defs := find_definitions(statement, name):
+                return defs[0]
         return None
 
     def verbosity(self) -> int:

mypy/plugins/functools.py

@@ -307,7 +307,7 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
 
     # Technically, we should set definition to None here, since it will not be recovered
     # on warm cache runs in fixup.py. This however may hide some helpful info in error
-    # messages, so we are keeping it for now.
+    # messages, so we are keeping it for now. See also issue #20640.
     partially_applied = fn_type.copy_modified(
         arg_types=partial_types,
         arg_kinds=partial_kinds,

mypy/traverser.py

@@ -66,6 +66,7 @@
     SetExpr,
     SliceExpr,
     StarExpr,
+    Statement,
     StrExpr,
     SuperExpr,
     TempNode,
@@ -96,7 +97,7 @@
     StarredPattern,
     ValuePattern,
 )
-from mypy.visitor import NodeVisitor
+from mypy.visitor import NodeVisitor, StatementVisitor
 
 
 @trait
@@ -108,6 +109,10 @@ class TraverserVisitor(NodeVisitor[None]):
     should override visit methods to perform actions during
     traversal. Calling the superclass method allows reusing the
     traversal implementation.
+
+    TODO: split this into more limited visitor (e.g. statements-only etc).
+    This will improve performance since in many cases we don't need to recurse
+    all the way down in various visitors that subclass this.
     """
 
     def __init__(self) -> None:
@@ -1084,3 +1089,115 @@ def visit_assignment_stmt(self, stmt: AssignmentStmt) -> None:
 
     def visit_yield_from_expr(self, expr: YieldFromExpr) -> None:
         self.yield_from_expressions.append((expr, self.in_assignment))
+
+
+def find_definitions(o: Statement, name: str) -> list[Statement]:
+    visitor = DefinitionSeeker(name)
+    o.accept(visitor)
+    return visitor.found
+
+
+class DefinitionSeeker(StatementVisitor[None]):
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.found: list[Statement] = []
+
+    def visit_assignment_stmt(self, o: AssignmentStmt, /) -> None:
+        # TODO: support more kinds of locations (like assignment statements).
+        # the latter will be helpful for type old-style aliases.
+        pass
+
+    def visit_for_stmt(self, o: ForStmt, /) -> None:
+        o.body.accept(self)
+        if o.else_body:
+            o.else_body.accept(self)
+
+    def visit_with_stmt(self, o: WithStmt, /) -> None:
+        o.body.accept(self)
+
+    def visit_del_stmt(self, o: DelStmt, /) -> None:
+        pass
+
+    def visit_func_def(self, o: FuncDef, /) -> None:
+        if o.name == self.name:
+            self.found.append(o)
+
+    def visit_overloaded_func_def(self, o: OverloadedFuncDef, /) -> None:
+        if o.name == self.name:
+            self.found.append(o)
+
+    def visit_class_def(self, o: ClassDef, /) -> None:
+        if o.name == self.name:
+            self.found.append(o)
+
+    def visit_global_decl(self, o: GlobalDecl, /) -> None:
+        pass
+
+    def visit_nonlocal_decl(self, o: NonlocalDecl, /) -> None:
+        pass
+
+    def visit_decorator(self, o: Decorator, /) -> None:
+        if o.name == self.name:
+            self.found.append(o)
+
+    def visit_import(self, o: Import, /) -> None:
+        pass
+
+    def visit_import_from(self, o: ImportFrom, /) -> None:
+        pass
+
+    def visit_import_all(self, o: ImportAll, /) -> None:
+        pass
+
+    def visit_block(self, o: Block, /) -> None:
+        for s in o.body:
+            s.accept(self)
+
+    def visit_expression_stmt(self, o: ExpressionStmt, /) -> None:
+        pass
+
+    def visit_operator_assignment_stmt(self, o: OperatorAssignmentStmt, /) -> None:
+        pass
+
+    def visit_while_stmt(self, o: WhileStmt, /) -> None:
+        o.body.accept(self)
+        if o.else_body:
+            o.else_body.accept(self)
+
+    def visit_return_stmt(self, o: ReturnStmt, /) -> None:
+        pass
+
+    def visit_assert_stmt(self, o: AssertStmt, /) -> None:
+        pass
+
+    def visit_if_stmt(self, o: IfStmt, /) -> None:
+        for b in o.body:
+            b.accept(self)
+        if o.else_body:
+            o.else_body.accept(self)
+
+    def visit_break_stmt(self, o: BreakStmt, /) -> None:
+        pass
+
+    def visit_continue_stmt(self, o: ContinueStmt, /) -> None:
+        pass
+
+    def visit_pass_stmt(self, o: PassStmt, /) -> None:
+        pass
+
+    def visit_raise_stmt(self, o: RaiseStmt, /) -> None:
+        pass
+
+    def visit_try_stmt(self, o: TryStmt, /) -> None:
+        o.body.accept(self)
+        if o.else_body is not None:
+            o.else_body.accept(self)
+        if o.finally_body is not None:
+            o.finally_body.accept(self)
+
+    def visit_match_stmt(self, o: MatchStmt, /) -> None:
+        for b in o.bodies:
+            b.accept(self)
+
+    def visit_type_alias_stmt(self, o: TypeAliasStmt, /) -> None:
+        pass

test-data/unit/check-incremental.test

@@ -7730,3 +7730,22 @@ tmp/b.py:2: note: "lol" defined here
 [out4]
 tmp/a.py:2: error: Unexpected keyword argument "uhhhh" for "lol"
 tmp/b.py:2: note: "lol" defined here
+
+[case testCachedUnexpectedKeywordArgumentNested]
+import a
+[file a.py]
+import b
+b.lol(uhhhh=12) # tweak
+[file a.py.2]
+import b
+b.lol(uhhhh=12)
+[file b.py]
+while True:
+    if True:
+        def lol() -> None: pass
+[out]
+tmp/a.py:2: error: Unexpected keyword argument "uhhhh" for "lol"
+tmp/b.py:3: note: "lol" defined here
+[out2]
+tmp/a.py:2: error: Unexpected keyword argument "uhhhh" for "lol"
+tmp/b.py:3: note: "lol" defined here