Update processing __init__.pyi in QualifiedNameFinder.PathChoosingVisitor so it provides correct names (PY-30942)

python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassTypeProvider.kt

@@ -68,7 +68,7 @@ class PyDataclassTypeProvider : PyTypeProviderBase() {
     if (resolvedCallee is PyCallable) {
       val instanceName = when (resolvedCallee.qualifiedName) {
         "dataclasses.replace" -> "obj"
-        "attr.__init__.assoc", "attr.__init__.evolve" -> "inst"
+        "attr.assoc", "attr.evolve" -> "inst"
         else -> return null
       }
 

python/src/com/jetbrains/python/inspections/PyDataclassInspection.kt

@@ -25,12 +25,12 @@ class PyDataclassInspection : PyInspection() {
   companion object {
     private val ORDER_OPERATORS = setOf("__lt__", "__le__", "__gt__", "__ge__")
     private val DATACLASSES_HELPERS = setOf("dataclasses.fields", "dataclasses.asdict", "dataclasses.astuple", "dataclasses.replace")
-    private val ATTRS_HELPERS = setOf("attr.__init__.fields",
-                                      "attr.__init__.fields_dict",
-                                      "attr.__init__.asdict",
-                                      "attr.__init__.astuple",
-                                      "attr.__init__.assoc",
-                                      "attr.__init__.evolve")
+    private val ATTRS_HELPERS = setOf("attr.fields",
+                                      "attr.fields_dict",
+                                      "attr.asdict",
+                                      "attr.astuple",
+                                      "attr.assoc",
+                                      "attr.evolve")
   }
 
   override fun buildVisitor(holder: ProblemsHolder,
@@ -105,7 +105,8 @@ class PyDataclassInspection : PyInspection() {
             this::registerProblem,
             {
               val stub = it.stub
-              val fieldStub = if (stub == null) PyDataclassFieldStubImpl.create(it) else stub.getCustomStub(PyDataclassFieldStub::class.java)
+              val fieldStub = if (stub == null) PyDataclassFieldStubImpl.create(it)
+              else stub.getCustomStub(PyDataclassFieldStub::class.java)
 
               fieldStub?.initValue() != false && !PyTypingTypeProvider.isClassVar(it, myTypeEvalContext)
             },
@@ -555,11 +556,10 @@ class PyDataclassInspection : PyInspection() {
     private fun processHelperAttrsArgument(argument: PyExpression?, calleeQName: String) {
       if (argument == null) return
 
-      val instance = calleeQName != "attr.__init__.fields" && calleeQName != "attr.__init__.fields_dict"
+      val instance = calleeQName != "attr.fields" && calleeQName != "attr.fields_dict"
 
       if (isNotExpectedDataclass(myTypeEvalContext.getType(argument), PyDataclassParameters.Type.ATTRS, !instance, instance)) {
-        val presentableCalleeQName = calleeQName.replaceFirst(".__init__.", ".")
-        val message = "'$presentableCalleeQName' method should be called on attrs " + if (instance) "instances" else "types"
+        val message = "'$calleeQName' method should be called on attrs " + if (instance) "instances" else "types"
 
         registerProblem(argument, message)
       }

python/src/com/jetbrains/python/psi/PyKnownDecoratorUtil.java

@@ -72,10 +72,11 @@ public class PyKnownDecoratorUtil {
     KOMBU_UTILS_CACHED_PROPERTY("kombu.utils.cached_property"),
 
     DATACLASSES_DATACLASS("dataclasses.dataclass"),
-    ATTR_S("attr.__init__.s"),
-    ATTR_ATTRS("attr.__init__.attrs"),
-    ATTR_ATTRIBUTES("attr.__init__.attributes"),
-    ATTR_DATACLASS("attr.__init__.dataclass"),
+    ATTR_S("attr.s"),
+    ATTR_ATTRS("attr.attrs"),
+    ATTR_ATTRIBUTES("attr.attributes"),
+    ATTR_DATACLASS("attr.dataclass"),
+
     PYTEST_FIXTURE("pytest.fixture");
 
     private final QualifiedName myQualifiedName;

python/src/com/jetbrains/python/psi/resolve/QualifiedNameFinder.java

@@ -71,7 +71,7 @@ public class QualifiedNameFinder {
 
   @Nullable
   private static QualifiedName shortestQName(@NotNull List<QualifiedName> qNames) {
-    return qNames.stream().min((o1, o2) -> o1.getComponentCount() - o2.getComponentCount()).orElse(null);
+    return qNames.stream().min(Comparator.comparingInt(QualifiedName::getComponentCount)).orElse(null);
   }
 
   @Nullable
@@ -198,39 +198,49 @@ public class QualifiedNameFinder {
    * Tries to find roots that contain given vfile.
    */
   private static class PathChoosingVisitor implements RootVisitor {
-    @Nullable private final VirtualFile myVFile;
+    @NotNull private final VirtualFile myVFile;
     @NotNull private final Set<QualifiedName> myResults = new LinkedHashSet<>();
 
     private PathChoosingVisitor(@NotNull VirtualFile file) {
-      if (!file.isDirectory() && file.getName().equals(PyNames.INIT_DOT_PY)) {
-        myVFile = file.getParent();
-      }
-      else {
-        myVFile = file;
-      }
+      myVFile = file;
     }
 
     @Override
     public boolean visitRoot(@NotNull VirtualFile root, @Nullable Module module, @Nullable Sdk sdk, boolean isModuleSource) {
-      if (myVFile != null) {
-        final String relativePath = VfsUtilCore.getRelativePath(myVFile, root, '/');
-        if (!StringUtil.isEmpty(relativePath)) {
-          final List<String> result = StringUtil.split(relativePath, "/");
-          if (!result.isEmpty()) {
-            final int lastIndex = result.size() - 1;
-            result.set(lastIndex, FileUtil.getNameWithoutExtension(result.get(lastIndex)));
+      final List<String> result = pathToNameComponents(VfsUtilCore.getRelativePath(myVFile, root, '/'));
+      if (!result.isEmpty()) {
+        for (String component : result) {
+          if (!PyNames.isIdentifier(component)) {
+            return true;
           }
-          for (String component : result) {
-            if (!PyNames.isIdentifier(component)) {
-              return true;
-            }
-          }
-          myResults.add(QualifiedName.fromComponents(result));
         }
+        myResults.add(QualifiedName.fromComponents(result));
       }
       return true;
     }
 
+    @NotNull
+    private List<String> pathToNameComponents(@Nullable String relativePath) {
+      if (StringUtil.isEmpty(relativePath)) return Collections.emptyList();
+
+      final List<String> result = new ArrayList<>(StringUtil.split(relativePath, "/"));
+      if (!result.isEmpty()) {
+        final int lastIndex = result.size() - 1;
+        final String nameWithoutExtension = FileUtil.getNameWithoutExtension(result.get(lastIndex));
+
+        if (myVFile.isDirectory() || !nameWithoutExtension.equals(PyNames.INIT)) {
+          result.set(lastIndex, nameWithoutExtension);
+        }
+        else {
+          result.remove(lastIndex);
+        }
+
+        return result;
+      }
+
+      return Collections.emptyList();
+    }
+
     @NotNull
     public List<QualifiedName> getResults() {
       return new ArrayList<>(myResults);