PY-33802 Support create a class quick fix for capitalized undefined references

GitOrigin-RevId: 975454ad1a7013a19f04be6121a500e80fb71075

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);
   }
 
 }

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();
           }

python/testData/inspections/addClassToImportedModule/addClassToImportedModule.py

@@ -0,0 +1,3 @@
+import mod
+
+clzz = mod.<caret><warning descr="Cannot find reference 'Clzz' in 'mod.py'">Clzz</warning>()

python/testData/inspections/addClassToImportedModule/mod_after.py

@@ -0,0 +1,2 @@
+class Clzz(object):
+    pass

python/testData/inspections/addClassToImportedPackage/addClassToImportedPackage.py

@@ -0,0 +1,3 @@
+import pkg
+
+clzz = pkg.<caret><warning descr="Cannot find reference 'Clzz' in '__init__.py'">Clzz</warning>()

python/testData/inspections/addClassToImportedPackage/pkg/__init___after.py

@@ -0,0 +1,2 @@
+class Clzz(object):
+    pass

python/testData/inspections/addClassToModuleInFromImport/addClassToModuleInFromImport.py

@@ -0,0 +1,4 @@
+from mod import <error descr="Cannot find reference 'Clzz' in 'mod.py'"><caret>Clzz</error>
+
+
+clzz = Clzz()

python/testData/inspections/addClassToModuleInFromImport/mod_after.py

@@ -0,0 +1,2 @@
+class Clzz(object):
+    pass

python/testData/inspections/addClassToPackageInFromImport/addClassToPackageInFromImport.py

@@ -0,0 +1,3 @@
+from mypack import <error descr="Cannot find reference 'Clzz' in '__init__.py'"><caret>Clzz</error>
+
+clzz = Clzz()

python/testData/inspections/addClassToPackageInFromImport/mypack/__init___after.py

@@ -0,0 +1,2 @@
+class Clzz(object):
+    pass

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,