From 19b0affab555b8a49cb95173cff1ecf57ee49716 Mon Sep 17 00:00:00 2001 From: Andrey Matveev Date: Thu, 17 Dec 2020 12:46:21 +0000 Subject: [PATCH] PY-33802 Support create a class quick fix for capitalized undefined references GitOrigin-RevId: 975454ad1a7013a19f04be6121a500e80fb71075 --- .../quickfix/CreateClassQuickFix.java | 6 ++--- .../PyUnresolvedReferencesVisitor.java | 22 ++++++++++------ .../addClassToImportedModule.py | 3 +++ .../addClassToImportedModule/mod.py | 0 .../addClassToImportedModule/mod_after.py | 2 ++ .../addClassToImportedPackage.py | 3 +++ .../addClassToImportedPackage/pkg/__init__.py | 0 .../pkg/__init___after.py | 2 ++ .../addClassToModuleInFromImport.py | 4 +++ .../addClassToModuleInFromImport/mod.py | 0 .../addClassToModuleInFromImport/mod_after.py | 2 ++ .../addClassToPackageInFromImport.py | 3 +++ .../mypack/__init__.py | 0 .../mypack/__init___after.py | 2 ++ .../com/jetbrains/python/PyQuickFixTest.java | 25 +++++++++++++++++++ 15 files changed, 64 insertions(+), 10 deletions(-) create mode 100644 python/testData/inspections/addClassToImportedModule/addClassToImportedModule.py create mode 100644 python/testData/inspections/addClassToImportedModule/mod.py create mode 100644 python/testData/inspections/addClassToImportedModule/mod_after.py create mode 100644 python/testData/inspections/addClassToImportedPackage/addClassToImportedPackage.py create mode 100644 python/testData/inspections/addClassToImportedPackage/pkg/__init__.py create mode 100644 python/testData/inspections/addClassToImportedPackage/pkg/__init___after.py create mode 100644 python/testData/inspections/addClassToModuleInFromImport/addClassToModuleInFromImport.py create mode 100644 python/testData/inspections/addClassToModuleInFromImport/mod.py create mode 100644 python/testData/inspections/addClassToModuleInFromImport/mod_after.py create mode 100644 python/testData/inspections/addClassToPackageInFromImport/addClassToPackageInFromImport.py create mode 100644 python/testData/inspections/addClassToPackageInFromImport/mypack/__init__.py create mode 100644 python/testData/inspections/addClassToPackageInFromImport/mypack/__init___after.py diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/CreateClassQuickFix.java b/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/CreateClassQuickFix.java index cd7e1e4d8c74..4fd1ac7edb53 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/CreateClassQuickFix.java +++ b/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/CreateClassQuickFix.java @@ -42,8 +42,8 @@ public class CreateClassQuickFix implements LocalQuickFix { @Override @NotNull public String getName() { - if (myAnchor instanceof PyFile) { - return PyPsiBundle.message("QFIX.create.class.in.module", myClassName, ((PyFile)myAnchor).getName()); + if (myAnchor.getElement() instanceof PyFile) { + return PyPsiBundle.message("QFIX.create.class.in.module", myClassName, ((PyFile)myAnchor.getElement()).getName()); } return PyPsiBundle.message("QFIX.create.class.0", myClassName); } @@ -79,7 +79,7 @@ public class CreateClassQuickFix implements LocalQuickFix { builder.replaceElement(pyClass.getSuperClassExpressions() [0], "object"); builder.replaceElement(pyClass.getStatementList(), PyNames.PASS); - PythonTemplateRunner.runTemplateInSelectedEditor(project, anchor, builder); + PythonTemplateRunner.runTemplate(pyClass.getContainingFile(), builder); } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/unresolvedReference/PyUnresolvedReferencesVisitor.java b/python/python-psi-impl/src/com/jetbrains/python/inspections/unresolvedReference/PyUnresolvedReferencesVisitor.java index b7a747c5f88a..adbfe008fd03 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/inspections/unresolvedReference/PyUnresolvedReferencesVisitor.java +++ b/python/python-psi-impl/src/com/jetbrains/python/inspections/unresolvedReference/PyUnresolvedReferencesVisitor.java @@ -972,8 +972,13 @@ public abstract class PyUnresolvedReferencesVisitor extends PyInspectionVisitor } else if (type instanceof PyModuleType) { PyFile file = ((PyModuleType)type).getModule(); - result.add(new AddFunctionQuickFix(refText, file.getName())); - getCreateClassFix(refText, element); + LocalQuickFix createClassQuickFix = getCreateClassFix(refText, element); + if (createClassQuickFix != null) { + result.add(createClassQuickFix); + } + else { + result.add(new AddFunctionQuickFix(refText, file.getName())); + } } return result; } @@ -1032,13 +1037,16 @@ public abstract class PyUnresolvedReferencesVisitor extends PyInspectionVisitor } LocalQuickFix getCreateClassFix(@NonNls String refText, PsiElement element) { - if (refText.length() > 2 && Character.isUpperCase(refText.charAt(0)) && !StringUtil.toUpperCase(refText).equals(refText) && - PsiTreeUtil.getParentOfType(element, PyImportStatementBase.class) == null) { + if (refText.length() > 2 && Character.isUpperCase(refText.charAt(0)) && !StringUtil.toUpperCase(refText).equals(refText)) { PsiElement anchor = element; if (element instanceof PyQualifiedExpression) { - final PyExpression expr = ((PyQualifiedExpression)element).getQualifier(); - if (expr != null) { - final PyType type = myTypeEvalContext.getType(expr); + PyExpression qualifier = ((PyQualifiedExpression)element).getQualifier(); + if (qualifier == null) { + final PyFromImportStatement fromImport = PsiTreeUtil.getParentOfType(element, PyFromImportStatement.class); + if (fromImport != null) qualifier = fromImport.getImportSource(); + } + if (qualifier != null) { + final PyType type = myTypeEvalContext.getType(qualifier); if (type instanceof PyModuleType) { anchor = ((PyModuleType)type).getModule(); } diff --git a/python/testData/inspections/addClassToImportedModule/addClassToImportedModule.py b/python/testData/inspections/addClassToImportedModule/addClassToImportedModule.py new file mode 100644 index 000000000000..51da5701a9a8 --- /dev/null +++ b/python/testData/inspections/addClassToImportedModule/addClassToImportedModule.py @@ -0,0 +1,3 @@ +import mod + +clzz = mod.Clzz() diff --git a/python/testData/inspections/addClassToImportedModule/mod.py b/python/testData/inspections/addClassToImportedModule/mod.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/inspections/addClassToImportedModule/mod_after.py b/python/testData/inspections/addClassToImportedModule/mod_after.py new file mode 100644 index 000000000000..d534b550ecb0 --- /dev/null +++ b/python/testData/inspections/addClassToImportedModule/mod_after.py @@ -0,0 +1,2 @@ +class Clzz(object): + pass \ No newline at end of file diff --git a/python/testData/inspections/addClassToImportedPackage/addClassToImportedPackage.py b/python/testData/inspections/addClassToImportedPackage/addClassToImportedPackage.py new file mode 100644 index 000000000000..f69d9a1ca18d --- /dev/null +++ b/python/testData/inspections/addClassToImportedPackage/addClassToImportedPackage.py @@ -0,0 +1,3 @@ +import pkg + +clzz = pkg.Clzz() diff --git a/python/testData/inspections/addClassToImportedPackage/pkg/__init__.py b/python/testData/inspections/addClassToImportedPackage/pkg/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/inspections/addClassToImportedPackage/pkg/__init___after.py b/python/testData/inspections/addClassToImportedPackage/pkg/__init___after.py new file mode 100644 index 000000000000..d534b550ecb0 --- /dev/null +++ b/python/testData/inspections/addClassToImportedPackage/pkg/__init___after.py @@ -0,0 +1,2 @@ +class Clzz(object): + pass \ No newline at end of file diff --git a/python/testData/inspections/addClassToModuleInFromImport/addClassToModuleInFromImport.py b/python/testData/inspections/addClassToModuleInFromImport/addClassToModuleInFromImport.py new file mode 100644 index 000000000000..bb8badfa00e3 --- /dev/null +++ b/python/testData/inspections/addClassToModuleInFromImport/addClassToModuleInFromImport.py @@ -0,0 +1,4 @@ +from mod import Clzz + + +clzz = Clzz() diff --git a/python/testData/inspections/addClassToModuleInFromImport/mod.py b/python/testData/inspections/addClassToModuleInFromImport/mod.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/inspections/addClassToModuleInFromImport/mod_after.py b/python/testData/inspections/addClassToModuleInFromImport/mod_after.py new file mode 100644 index 000000000000..d534b550ecb0 --- /dev/null +++ b/python/testData/inspections/addClassToModuleInFromImport/mod_after.py @@ -0,0 +1,2 @@ +class Clzz(object): + pass \ No newline at end of file diff --git a/python/testData/inspections/addClassToPackageInFromImport/addClassToPackageInFromImport.py b/python/testData/inspections/addClassToPackageInFromImport/addClassToPackageInFromImport.py new file mode 100644 index 000000000000..082fa1c3105e --- /dev/null +++ b/python/testData/inspections/addClassToPackageInFromImport/addClassToPackageInFromImport.py @@ -0,0 +1,3 @@ +from mypack import Clzz + +clzz = Clzz() \ No newline at end of file diff --git a/python/testData/inspections/addClassToPackageInFromImport/mypack/__init__.py b/python/testData/inspections/addClassToPackageInFromImport/mypack/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/inspections/addClassToPackageInFromImport/mypack/__init___after.py b/python/testData/inspections/addClassToPackageInFromImport/mypack/__init___after.py new file mode 100644 index 000000000000..d534b550ecb0 --- /dev/null +++ b/python/testData/inspections/addClassToPackageInFromImport/mypack/__init___after.py @@ -0,0 +1,2 @@ +class Clzz(object): + pass \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java index 65ae3613937b..27b3c86ddbb3 100644 --- a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java @@ -222,6 +222,31 @@ public class PyQuickFixTest extends PyTestCase { () -> doInspectionTest(PyUnresolvedReferencesInspection.class, "Create class 'MyClass'", true, true)); } + // PY-33802 + public void testAddClassToImportedModule() { + doMultiFilesInspectionTest(PyUnresolvedReferencesInspection.class, + PyPsiBundle.message("QFIX.create.class.in.module", "Clzz", "mod.py"), "mod.py"); + } + + // PY-33802 + public void testAddClassToImportedPackage() { + doMultiFilesInspectionTest(PyUnresolvedReferencesInspection.class, + PyPsiBundle.message("QFIX.create.class.in.module", "Clzz", "__init__.py"), "pkg/__init__.py"); + } + + // PY-33802 + public void testAddClassToModuleInFromImport() { + doMultiFilesInspectionTest(PyUnresolvedReferencesInspection.class, + PyPsiBundle.message("QFIX.create.class.in.module", "Clzz", "mod.py"), "mod.py"); + } + + // PY-33802 + public void testAddClassToPackageInFromImport() { + doMultiFilesInspectionTest(PyUnresolvedReferencesInspection.class, + PyPsiBundle.message("QFIX.create.class.in.module", "Clzz", "__init__.py"), + "mypack/__init__.py"); + } + // PY-21204 public void testAddFunctionFromFString() { runWithLanguageLevel(LanguageLevel.PYTHON36,