PY-45588 Generalize PyOverrideImplementUtil.addImports to cover parameter defaults

GitOrigin-RevId: 828305e2a7449704c08503d8d68394d28ce98c83

python/src/com/jetbrains/python/codeInsight/override/PyOverrideImplementUtil.java

@@ -169,14 +169,14 @@ public final class PyOverrideImplementUtil {
     final TypeEvalContext context = TypeEvalContext.userInitiated(cls.getProject(), cls.getContainingFile());
 
     final PyFunction function = buildOverriddenFunction(cls, baseFunction, implement).addFunctionAfter(statementList, anchor);
-    addImports(baseFunction, function, context);
+    addImports(baseFunction, function);
 
     PyiUtil
       .getOverloads(baseFunction, context)
       .forEach(
         baseOverload -> {
           final PyFunction overload = (PyFunction)statementList.addBefore(baseOverload, function);
-          addImports(baseOverload, overload, context);
+          addImports(baseOverload, overload);
         }
       );
 
@@ -336,36 +336,20 @@ public final class PyOverrideImplementUtil {
    * @param baseFunction base function used to resolve types
    * @param function overridden function
    */
-  private static void addImports(@NotNull PyFunction baseFunction, @NotNull PyFunction function, @NotNull TypeEvalContext context) {
-    final UnresolvedExpressionVisitor unresolvedExpressionVisitor = new UnresolvedExpressionVisitor();
-    getAnnotations(function, context).forEach(annotation -> annotation.accept(unresolvedExpressionVisitor));
-    getDecorators(function).forEach(decorator -> decorator.accept(unresolvedExpressionVisitor));
+  private static void addImports(@NotNull PyFunction baseFunction, @NotNull PyFunction function) {
+    UnresolvedExpressionVisitor unresolvedExpressionVisitor = new UnresolvedExpressionVisitor() {
+      @Override
+      public void visitPyStatementList(@NotNull PyStatementList node) {
+      }
+    };
+    function.accept(unresolvedExpressionVisitor);
 
-    final ResolveExpressionVisitor resolveExpressionVisitor = new ResolveExpressionVisitor(unresolvedExpressionVisitor.getUnresolved());
-    getAnnotations(baseFunction, context).forEach(annotation -> annotation.accept(resolveExpressionVisitor));
-    getDecorators(baseFunction).forEach(decorator -> decorator.accept(resolveExpressionVisitor));
-  }
-
-  /**
-   * Collect annotations from function parameters and return.
-   *
-   */
-  @NotNull
-  private static List<PyAnnotation> getAnnotations(@NotNull PyFunction function, @NotNull TypeEvalContext typeEvalContext) {
-    return StreamEx.of(function.getParameters(typeEvalContext))
-        .map(PyCallableParameter::getParameter)
-        .select(PyNamedParameter.class)
-        .remove(PyParameter::isSelf)
-        .map(PyAnnotationOwner::getAnnotation)
-        .append(function.getAnnotation())
-        .nonNull()
-        .toList();
-  }
-
-  @NotNull
-  private static List<PyDecorator> getDecorators(@NotNull PyFunction function) {
-    final PyDecoratorList decoratorList = function.getDecoratorList();
-    return decoratorList == null ? Collections.emptyList() : Arrays.asList(decoratorList.getDecorators());
+    ResolveExpressionVisitor resolveExpressionVisitor = new ResolveExpressionVisitor(unresolvedExpressionVisitor.getUnresolved()) {
+      @Override
+      public void visitPyStatementList(@NotNull PyStatementList node) {
+      }
+    };
+    baseFunction.accept(resolveExpressionVisitor);
   }
 
   /**

python/testData/override/ImportForParameterDefaultValue/main.py

@@ -0,0 +1,5 @@
+from mod import Base
+
+
+class Sub(Base):
+    pass

python/testData/override/ImportForParameterDefaultValue/main_after.py

@@ -0,0 +1,6 @@
+from mod import Base, default
+
+
+class Sub(Base):
+    def method(self, param=default):
+        super().method(param)

python/testData/override/ImportForParameterDefaultValue/mod.py

@@ -0,0 +1,6 @@
+default = object()
+
+
+class Base:
+    def method(self, param=default):
+        pass

python/testSrc/com/jetbrains/python/PyOverrideTest.java

@@ -244,6 +244,13 @@ public class PyOverrideTest extends PyTestCase {
     });
   }
 
+  public void testImportForParameterDefaultValue() {
+    myFixture.copyDirectoryToProject(getTestName(false), "");
+    myFixture.configureByFile("main.py");
+    doOverride(null);
+    myFixture.checkResultByFile(getTestName(false) + "/main_after.py", true);
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/override";