PY-54167 Add import statement for class when using completion in

__all__

GitOrigin-RevId: 4e54ad17a5965e80557e822ae6a9d323da428f9f
This commit is contained in:
chbndrhnns
2025-03-14 10:39:56 +01:00
committed by intellij-monorepo-bot
parent 7e0472a5c7
commit 91bddd1e62
2 changed files with 39 additions and 0 deletions

View File

@@ -710,6 +710,13 @@ public abstract class PythonCommonCompletionTest extends PythonCommonTestCase {
doTest();
}
public void testDunderAllClassReference() { // PY-54167
myFixture.copyDirectoryToProject(getTestName(true), "");
myFixture.configureByFile("__init__.py");
myFixture.complete(CompletionType.BASIC, 2);
myFixture.checkResultByFile(getTestName(true) + "/__init__.after.py");
}
public void testDunderAllReferenceImport() { // PY-6306
doTest();
}

View File

@@ -279,6 +279,9 @@ public final class PyClassNameCompletionContributor extends PyImportableNameComp
@NotNull PsiElement position,
@NotNull TypeEvalContext typeEvalContext) {
if (position.getParent() instanceof PyStringLiteralExpression) {
if (isInsideAllInInitPy(position)) {
return getImportingInsertHandler();
}
return getStringLiteralInsertHandler();
}
// Some names in typing are defined as functions, this rule needs to have priority
@@ -291,6 +294,35 @@ public final class PyClassNameCompletionContributor extends PyImportableNameComp
return getImportingInsertHandler();
}
private static boolean isInsideAllInInitPy(@NotNull PsiElement position) {
PsiFile originalFile = position.getContainingFile();
if (originalFile == null) {
return false;
}
if (!PyNames.INIT_DOT_PY.equals(originalFile.getName())) {
return false;
}
PyAssignmentStatement assignment =
PsiTreeUtil.getParentOfType(position, PyAssignmentStatement.class);
if (assignment == null) {
return false;
}
PyExpression[] targets = assignment.getTargets();
if (targets.length != 1) {
return false;
}
PyExpression target = targets[0];
if (!(target instanceof PyTargetExpression firstTarget)) {
return false;
}
return PyNames.ALL.equals(firstTarget.getName());
}
private static class Counters {
int scannedNames;
int privateNames;