PY-73246 Automatically insert square brackets on completion of parameterized types

GitOrigin-RevId: 5e81bc984fa3c7b4f83c564c85d872c5a49cf391
2026-09-27 10:03:11 +07:00 · 2024-07-02 14:08:48 +00:00
parent 52850e21d8
commit 1a3e6c2a64
26 changed files with 141 additions and 9 deletions
@@ -132,7 +132,7 @@ public final class PyClassNameCompletionContributor extends PyImportableNameComp
                  presentation.setTypeText(importPath.toString());
                }
              })
-              .withInsertHandler(getInsertHandler(exported, position));
+              .withInsertHandler(getInsertHandler(exported, position, typeEvalContext));
            result.addElement(PrioritizedLookupElement.withPriority(lookupElement, PythonCompletionWeigher.NOT_IMPORTED_MODULE_WEIGHT));
            counters.totalVariants++;
            if (counters.totalVariants >= maxVariants) return false;
@@ -211,10 +211,15 @@ public final class PyClassNameCompletionContributor extends PyImportableNameComp
  }

  private @NotNull InsertHandler<LookupElement> getInsertHandler(@NotNull PyElement exported,
-                                                                 @NotNull PsiElement position) {
+                                                                 @NotNull PsiElement position,
+                                                                 @NotNull TypeEvalContext typeEvalContext) {
    if (position.getParent() instanceof PyStringLiteralExpression) {
      return getStringLiteralInsertHandler();
    }
+    // Some names in typing are defined as functions, this rule needs to have priority
+    else if (PyParameterizedTypeInsertHandler.isCompletingParameterizedType(exported, position, typeEvalContext)) {
+      return getGenericTypeInsertHandler();
+    }
    else if (exported instanceof PyFunction && !(position.getParent().getParent() instanceof PyDecorator)) {
      return getFunctionInsertHandler();
    }
@@ -30,6 +30,13 @@ abstract class PyImportableNameCompletionContributor : CompletionContributor(),
      addImportForLookupElement(context, item, tailOffset)
    }
  }
+  
+  protected val genericTypeInsertHandler: InsertHandler<LookupElement> = InsertHandler<LookupElement> { context, item ->
+    val tailOffset = context.tailOffset - 1
+    PyParameterizedTypeInsertHandler.INSTANCE.handleInsert(context, item)
+    context.commitDocument()
+    addImportForLookupElement(context, item, tailOffset)
+  } 

  protected val stringLiteralInsertHandler: InsertHandler<LookupElement> = InsertHandler { context, item ->
    val element = item.psiElement
@@ -0,0 +1,40 @@
+package com.jetbrains.python.codeInsight.completion;
+
+import com.intellij.codeInsight.completion.InsertionContext;
+import com.intellij.codeInsight.completion.util.ParenthesesInsertHandler;
+import com.intellij.codeInsight.lookup.LookupElement;
+import com.intellij.psi.PsiElement;
+import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider;
+import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.PyQualifiedNameOwner;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Inserts square brackets after names of parameterized types inside type hints.
+ * It includes special constructs from typings, such as {@code typing.Optional} or {@code typing.Union},
+ * and regular generic classes, such as {@code list} or {@code ContextManager}.
+ */
+public class PyParameterizedTypeInsertHandler extends ParenthesesInsertHandler<LookupElement> {
+  public static boolean isCompletingParameterizedType(@NotNull PsiElement definition,
+                                                      @NotNull PsiElement completionPosition,
+                                                      @NotNull TypeEvalContext context) {
+    return PyTypingTypeProvider.isInsideTypeHint(completionPosition, context) && isParameterizedType(definition, context);
+  }
+
+  public static final PyParameterizedTypeInsertHandler INSTANCE = new PyParameterizedTypeInsertHandler();
+
+  private PyParameterizedTypeInsertHandler() {
+    super(false, false, true, false, '[', ']');
+  }
+
+  @Override
+  protected boolean placeCaretInsideParentheses(InsertionContext context, LookupElement item) {
+    return true;
+  }
+
+  private static boolean isParameterizedType(@NotNull PsiElement element, @NotNull TypeEvalContext typeEvalContext) {
+    return element instanceof PyQualifiedNameOwner qnOwner && PyTypingTypeProvider.GENERIC_CLASSES.contains(qnOwner.getQualifiedName()) ||
+           element instanceof PyClass pyClass && PyTypingTypeProvider.isGeneric(pyClass, typeEvalContext);
+  }
+}
@@ -2021,6 +2021,11 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
      return true;
    }

+    PyTypeAliasStatement typeAlias = PsiTreeUtil.getParentOfType(realContext, PyTypeAliasStatement.class, false, PyStatement.class);
+    if (typeAlias != null && PsiTreeUtil.isAncestor(typeAlias.getTypeExpression(), realContext, false)) {
+      return true;
+    }
+
    return false;
  }

@@ -16,6 +16,7 @@ import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.completion.PyClassInsertHandler;
 import com.jetbrains.python.codeInsight.completion.PyFunctionInsertHandler;
+import com.jetbrains.python.codeInsight.completion.PyParameterizedTypeInsertHandler;
 import com.jetbrains.python.codeInsight.completion.PythonCompletionWeigher;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
@@ -67,10 +68,13 @@ public class CompletionVariantsProcessor extends VariantsProcessor {
      final Project project = element.getProject();
      final TypeEvalContext context = TypeEvalContext.codeCompletion(project, myContext != null ? myContext.getContainingFile() : null);

-      if (!mySuppressParentheses &&
-          element instanceof PyFunction && ((PyFunction)element).getProperty() == null &&
-          !PyKnownDecoratorUtil.hasUnknownDecorator((PyFunction)element, context) &&
-          !isSingleArgDecoratorCall(myContext, (PyFunction)element)) {
+      if (myContext != null && PyParameterizedTypeInsertHandler.isCompletingParameterizedType(element, myContext, context)) {
+        item = item.withInsertHandler(PyParameterizedTypeInsertHandler.INSTANCE);        
+      }
+      else if (!mySuppressParentheses &&
+               element instanceof PyFunction && ((PyFunction)element).getProperty() == null &&
+               !PyKnownDecoratorUtil.hasUnknownDecorator((PyFunction)element, context) &&
+               !isSingleArgDecoratorCall(myContext, (PyFunction)element)) {
        item = item.withInsertHandler(PyFunctionInsertHandler.INSTANCE);
        final List<PyCallableParameter> parameters = ((PyFunction)element).getParameters(context);
        final String params = StringUtil.join(parameters, PyCallableParameter::getName, ", ");
@@ -1,3 +1,3 @@
 from typing import List

-xs = []  # type: List
+xs = []  # type: List[]
@@ -2,4 +2,4 @@ from typing import Final


 class C:
-    attr: Final(<caret>)
+    attr: Final[<caret>]
@@ -1,3 +1,3 @@
 from typing import Tuple

-attr: Tuple<caret>
+attr: Tuple[<caret>]
@@ -0,0 +1 @@
+type Alias = list[<caret>]
@@ -0,0 +1 @@
+type Alias = lis<caret>
@@ -0,0 +1,4 @@
+from typing import TypeAlias
+
+
+Alias: TypeAlias = list[<caret>]
@@ -0,0 +1,4 @@
+from typing import TypeAlias
+
+
+Alias: TypeAlias = lis<caret>
@@ -0,0 +1 @@
+xs = []  # type: list[<caret>]
@@ -0,0 +1 @@
+xs = []  # type: lis<caret>
@@ -0,0 +1 @@
+xs: list[<caret>]
@@ -0,0 +1 @@
+xs: lis<caret>
@@ -0,0 +1,3 @@
+from mod import UniqueBox
+
+x: UniqueBox[<caret>]
@@ -0,0 +1 @@
+x: UniqueBo<caret>
@@ -0,0 +1,2 @@
+class UniqueBox[T]:
+    pass
@@ -0,0 +1,5 @@
+from typing import ClassVar
+
+
+class C:
+    attr: ClassVar[<caret>]
@@ -0,0 +1,2 @@
+class C:
+    attr: ClassV<caret>
@@ -0,0 +1,3 @@
+@_SpecialForm
+def ClassVar(self, parameters):
+    ...
@@ -0,0 +1,3 @@
+from mod import UniqueBox
+
+x = UniqueBox<caret>
@@ -0,0 +1 @@
+x = UniqueBo<caret>
@@ -0,0 +1,2 @@
+class UniqueBox[T]:
+    pass
@@ -708,6 +708,41 @@ public class Py3CompletionTest extends PyTestCase {
    assertContainsElements(suggested, "name=", "year=");
  }

+  // PY-73246
+  public void testSquareBracketsInsertedAfterImportedParameterizedTypesFromTypingInsideTypeHints() {
+    doMultiFileTest();
+  }
+  
+  // PY-73246
+  public void testSquareBracketsInsertedAfterImportedGenericClassesInsideTypeHints() {
+    doMultiFileTest();
+  }
+
+  // PY-73246
+  public void testSquareBracketsNotInsertedAfterImportedGenericClassesOutsideTypeHints() {
+    doMultiFileTest();
+  }
+
+  // PY-73246
+  public void testSquareBracketsInsertedAfterBuiltinGenericTypesInsideTypeHints() {
+    doMultiFileTest();
+  }
+
+  // PY-73246
+  public void testSquareBracketsInsertedAfterBuiltinGenericTypesInsideTypeComments() {
+    doMultiFileTest();
+  }
+  
+  // PY-73246
+  public void testSquareBracketsInsertedAfterBuiltinGenericTypesInsideOldStyleTypeAliases() {
+    doMultiFileTest();
+  }
+
+  // PY-73246
+  public void testSquareBracketsInsertedAfterBuiltinGenericTypesInsideNewStyleTypeAliases() {
+    doMultiFileTest();
+  }
+
  private void doTestVariants(String @NotNull ... expected) {
    final String testName = getTestName(true);
    myFixture.configureByFile(testName + ".py");