[python] PY-79283 add import for attribute completion

GitOrigin-RevId: 0c75055d515bce3e7bf5927df2ed90fa38668c14
This commit is contained in:
Morgan Bartholomew
2025-11-14 09:10:44 +00:00
committed by intellij-monorepo-bot
parent 71ae5608f6
commit 7bf4fae01e
6 changed files with 87 additions and 6 deletions
+1
View File
@@ -196,6 +196,7 @@ jvm_library(
"//python/impl.helperLocator:community-helpersLocator",
"//python/python-sdk:sdk",
"//python/python-sdk:sdk_test_lib",
"//platform/testFramework/core",
]
)
### auto-generated section `build intellij.python.commonTests` end
+1
View File
@@ -24,5 +24,6 @@
<orderEntry type="module" module-name="intellij.python.community" scope="TEST" />
<orderEntry type="module" module-name="intellij.python.community.helpersLocator" scope="TEST" />
<orderEntry type="module" module-name="intellij.python.sdk" scope="TEST" />
<orderEntry type="module" module-name="intellij.platform.testFramework.core" scope="TEST" />
</component>
</module>
@@ -446,6 +446,7 @@ class PyClassNameCompletionContributor : CompletionContributor(), DumbAware {
}
}
@JvmStatic
fun addImportForLookupElement(context: InsertionContext, item: LookupElement, tailOffset: Int) {
val manager = PsiDocumentManager.getInstance(context.project)
val document = manager.getDocument(context.file)
@@ -23,9 +23,11 @@ import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PythonCodeStyleService;
import com.jetbrains.python.ast.impl.PyUtilCore;
import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
import com.jetbrains.python.documentation.docstrings.DocStringUtil;
import com.jetbrains.python.documentation.doctest.PyDocstringFile;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyBuiltinCache;
import com.jetbrains.python.psi.impl.PyCodeFragmentWithHiddenImports;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
@@ -790,18 +792,28 @@ public final class AddImportHelper {
* @see #addOrUpdateFromImportStatement
*/
public static void addImport(@NotNull PsiNamedElement target, @NotNull PsiFile file, @NotNull PyElement element) {
if (target.getContainingFile().equals(file)) return;
if (PyBuiltinCache.getInstance(element).isBuiltin(target)) return;
if (target instanceof PsiFileSystemItem) {
addFileSystemItemImport((PsiFileSystemItem)target, file, element);
return;
}
final String name = target.getName();
// If target is a class attribute, import the containing class
PsiNamedElement elementToImport = target;
var parent = ScopeUtil.getScopeOwner(target);
if (parent instanceof PyClass pyClass) {
elementToImport = pyClass;
}
final String name = elementToImport.getName();
if (name == null) return;
final PsiFileSystemItem toImport = target.getContainingFile();
final PsiFileSystemItem toImport = elementToImport.getContainingFile();
if (toImport == null) return;
final QualifiedName importPath = QualifiedNameFinder.findCanonicalImportPath(target, element);
final QualifiedName importPath = QualifiedNameFinder.findCanonicalImportPath(elementToImport, element);
if (importPath == null) return;
final String path = importPath.toString();
@@ -811,7 +823,7 @@ public final class AddImportHelper {
addImportStatement(file, path, null, priority, element);
final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(file.getProject());
element.replace(elementGenerator.createExpressionFromText(LanguageLevel.forElement(target), path + "." + name));
element.replace(elementGenerator.createExpressionFromText(LanguageLevel.forElement(elementToImport), path + "." + name));
}
else {
addOrUpdateFromImportStatement(file, path, name, null, priority, element);
@@ -17,6 +17,8 @@ package com.jetbrains.python.psi.impl.references;
import com.intellij.codeInsight.completion.CompletionInitializationContext;
import com.intellij.codeInsight.completion.CompletionUtilCoreImpl;
import com.intellij.codeInsight.completion.InsertHandler;
import com.intellij.codeInsight.completion.InsertionContext;
import com.intellij.codeInsight.lookup.AutoCompletionPolicy;
import com.intellij.codeInsight.lookup.LookupElement;
import com.intellij.codeInsight.lookup.LookupElementBuilder;
@@ -28,7 +30,9 @@ import com.intellij.ui.IconManager;
import com.intellij.util.ArrayUtil;
import com.intellij.util.ObjectUtils;
import com.intellij.util.ProcessingContext;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.codeInsight.completion.PyClassNameCompletionContributor;
import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
import com.jetbrains.python.codeInsight.dataflow.scope.Scope;
@@ -213,7 +217,20 @@ public class PyQualifiedReference extends PyReferenceImpl {
else {
final PyClassType guessedType = guessClassTypeByName();
if (guessedType != null) {
Collections.addAll(variants, guessedType.getCompletionVariants(myElement.getName(), myElement, ctx));
ContainerUtil.addAll(variants, ContainerUtil.map(guessedType.getCompletionVariants(myElement.getName(), myElement, ctx), entry ->
entry instanceof LookupElementBuilder e
? e.withInsertHandler(new InsertHandler<LookupElement>() {
@Override
public void handleInsert(@NotNull InsertionContext context, @NotNull LookupElement item) {
PyClassNameCompletionContributor.InsertHandlers.addImportForLookupElement(context, item, context.getTailOffset() - 1);
if (e.getInsertHandler() != null) {
PsiDocumentManager.getInstance(context.getProject()).doPostponedOperationsAndUnblockDocument(context.getDocument());
e.handleInsert(context);
}
}
})
: element)
);
}
if (qualifier instanceof PyReferenceExpression) {
Collections.addAll(variants, collectSeenMembers(qualifier.getText(), ctx));
@@ -235,7 +252,7 @@ public class PyQualifiedReference extends PyReferenceImpl {
Collection<PyClass> classes = PyClassNameIndexInsensitive.find(className, getElement().getProject());
classes = filterByImports(classes, myElement.getContainingFile());
if (classes.size() == 1) {
return new PyClassTypeImpl(classes.iterator().next(), false);
return new PyClassTypeImpl(classes.iterator().next(), true);
}
}
}
@@ -6,6 +6,7 @@ import com.intellij.codeInsight.lookup.Lookup;
import com.intellij.codeInsight.lookup.LookupElement;
import com.intellij.codeInsight.lookup.LookupElementPresentation;
import com.intellij.codeInsight.lookup.LookupElementRenderer;
import com.intellij.idea.TestFor;
import com.intellij.openapi.module.Module;
import com.intellij.testFramework.PsiTestUtil;
import com.intellij.testFramework.TestDataPath;
@@ -898,6 +899,54 @@ public class Py3CompletionTest extends PyTestCase {
});
}
@TestFor(issues = "PY-79283")
public void testImportClassFromAttribute() {
myFixture.configureByText("mod.py", """
class Class:
unique_attribute = 1
def __init__(self):
self.unique_attribute_instance = 2
""");
doTestByText("Class.uniq<caret>");
myFixture.checkResult(
"""
from mod import Class
Class.unique_attribute"""
);
}
@TestFor(issues = "PY-79283")
public void testImportClassFromMethod() {
myFixture.configureByText("mod.py", """
class Class:
def unique_method(self): ...
""");
doTestByText("Class.uniq<caret>");
myFixture.checkResult(
"""
from mod import Class
Class.unique_method()"""
);
}
@TestFor(issues = "PY-79283")
public void testImportClassFromInnerClass() {
myFixture.configureByText("mod.py", """
class Class:
class UniqueClass: ...
""");
doTestByText("Class.Uniq<caret>");
myFixture.checkResult(
"""
from mod import Class
Class.UniqueClass"""
);
}
private void doTestVariants(String @NotNull ... expected) {
final String testName = getTestName(true);
myFixture.configureByFile(testName + ".py");