diff --git a/python/BUILD.bazel b/python/BUILD.bazel index 794372a5181f..6ded953675d6 100644 --- a/python/BUILD.bazel +++ b/python/BUILD.bazel @@ -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 diff --git a/python/intellij.python.commonTests.iml b/python/intellij.python.commonTests.iml index b30a1ea24307..cf614a7d234f 100644 --- a/python/intellij.python.commonTests.iml +++ b/python/intellij.python.commonTests.iml @@ -24,5 +24,6 @@ + \ No newline at end of file diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.kt index 6285b2b71d08..08b073ca20e5 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyClassNameCompletionContributor.kt @@ -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) diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java index b359e9120a0f..bb0e78bcb627 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java @@ -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); diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java index bcb5683205a3..d37013aaa797 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyQualifiedReference.java @@ -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() { + @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 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); } } } diff --git a/python/testSrc/com/jetbrains/python/Py3CompletionTest.java b/python/testSrc/com/jetbrains/python/Py3CompletionTest.java index bcee6b94ff72..2f476035216c 100644 --- a/python/testSrc/com/jetbrains/python/Py3CompletionTest.java +++ b/python/testSrc/com/jetbrains/python/Py3CompletionTest.java @@ -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"); + 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"); + 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"); + myFixture.checkResult( + """ + from mod import Class + + Class.UniqueClass""" + ); + } + private void doTestVariants(String @NotNull ... expected) { final String testName = getTestName(true); myFixture.configureByFile(testName + ".py");