[python] PY-77223 suggest more stubs

GitOrigin-RevId: 85cb25cd5286e8757e20a1e204740b80a6fdec84

python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyStubPackages.kt

@@ -19,6 +19,7 @@ import com.jetbrains.python.psi.resolve.RatedResolveResult
 import com.jetbrains.python.pyi.PyiFile
 import com.jetbrains.python.pyi.PyiUtil
 
+const val TYPES_PREFIX = "types-"
 const val STUBS_SUFFIX = "-stubs"
 private val STUB_PACKAGE_KEY = Key<Boolean>("PY_STUB_PACKAGE")
 private val INLINE_PACKAGE_KEY = Key<Boolean>("PY_INLINE_PACKAGE")
@@ -31,11 +32,17 @@ private val INLINE_PACKAGE_KEY = Key<Boolean>("PY_INLINE_PACKAGE")
 fun convertStubToRuntimePackageName(name: QualifiedName): QualifiedName {
   val top = name.firstComponent
 
-  if (top != null && top.endsWith(STUBS_SUFFIX)) {
-    return QualifiedName.fromComponents(top.removeSuffix(STUBS_SUFFIX)).append(name.removeHead(1))
+  return when {
+    top != null && top.endsWith(STUBS_SUFFIX) -> {
+      top.removeSuffix(STUBS_SUFFIX)
+    }
+    top != null && top.startsWith(TYPES_PREFIX) -> {
+      top.removePrefix(TYPES_PREFIX)
+    }
+    else -> return name
+  }.let {
+    QualifiedName.fromComponents(it).append(name.removeHead(1))
   }
-
-  return name
 }
 
 /**
@@ -56,6 +63,15 @@ fun findStubPackage(dir: PsiDirectory,
       doTransferStubPackageMarker(stubPackage)
       return stubPackage
     }
+
+    val typesPackageName = "$TYPES_PREFIX$referencedName"
+    val typesPackage = dir.findSubdirectory(typesPackageName)
+
+    // see comment about case sensitivity in com.jetbrains.python.psi.resolve.ResolveImportUtil.resolveInDirectory
+    if (typesPackage?.name == stubPackageName && (!checkForPackage || PyUtil.isPackage(typesPackage, dir))) {
+      doTransferStubPackageMarker(typesPackage)
+      return typesPackage
+    }
   }
 
   return null

python/python-psi-impl/src/com/jetbrains/python/packaging/PyRequirements.kt

@@ -20,6 +20,19 @@ import com.jetbrains.python.packaging.requirement.PyRequirementVersionSpec
  */
 fun pyRequirement(name: String): PyRequirement = PyRequirementImpl(name, emptyList(), listOf(name), "")
 
+/**
+ * This helper is not an API, consider using methods listed below.
+ *
+ * @see PyPackageManager.parseRequirement
+ * @see PyPackageManager.parseRequirements
+ *
+ * @see PyRequirementParser.fromLine
+ * @see PyRequirementParser.fromText
+ * @see PyRequirementParser.fromFile
+ */
+fun pyRequirement(name: String, relation: PyRequirementRelation, version: String): PyRequirement =
+  pyRequirement(name, relation, version, "")
+
 /**
  * This helper is not an API, consider using methods listed below.
  * If given version could not be normalized, then specified relation will be replaced with [PyRequirementRelation.STR_EQ].
@@ -33,9 +46,9 @@ fun pyRequirement(name: String): PyRequirement = PyRequirementImpl(name, emptyLi
  *
  * @see pyRequirementVersionSpec
  */
-fun pyRequirement(name: String, relation: PyRequirementRelation, version: String): PyRequirement {
+fun pyRequirement(name: String, relation: PyRequirementRelation, version: String, extras: String = ""): PyRequirement {
   val versionSpec = pyRequirementVersionSpec(relation, version)
-  return PyRequirementImpl(name, listOf(versionSpec), listOf(name + relation.presentableText + version), "")
+  return PyRequirementImpl(name, listOf(versionSpec), listOf(name + relation.presentableText + version), extras)
 }
 
 

python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesAdvertiser.kt

@@ -45,8 +45,15 @@ private class PyStubPackagesAdvertiser : PyInspection() {
                                 "pika" to "pika",
                                 "gi" to "PyGObject",
                                 "PyQt5" to "PyQt5",
+                                "pandas" to "pandas",
+                                "celery" to "celery",
+                                "urllib3" to "urllib3",
+                                "pillow" to "Pillow",
+                                "boto3" to "boto3",
                                 "traits" to "traits") // top-level package to package on PyPI, sorted by the latter
 
+    private val EXTRAS = mapOf("boto3-stubs" to "[full]")
+
     private val BALLOON_SHOWING = Key.create<Boolean>("showingStubPackagesAdvertiserBalloon")
   }
 
@@ -246,7 +253,7 @@ private class PyStubPackagesAdvertiser : PyInspection() {
         .flatMap { it.packages.entries.asSequence() }
         .filterNot { isIgnoredStubPackage(it.key, it.value.first, ignoredStubPackages) }
         .map {
-          pyRequirement(it.key, PyRequirementRelation.EQ, it.value.first)
+          pyRequirement(it.key, PyRequirementRelation.EQ, it.value.first, extras = EXTRAS.getOrDefault(it.key, ""))
         }
         .toList()
       if (requirements.isEmpty()) return emptyList<PyRequirement>() to emptyList()

python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesCompatibilityInspection.kt

@@ -36,8 +36,11 @@ class PyStubPackagesCompatibilityInspection : PyInspection() {
 
       return installedPackages
         .asSequence()
-        .filter { it.name.endsWith(STUBS_SUFFIX) && stubPkgsFilter(it) }
-        .mapNotNull { stubPkg -> nameToPkg[stubPkg.name.removeSuffix(STUBS_SUFFIX)]?.let { it to stubPkg } }
+        .filter { (it.name.endsWith(STUBS_SUFFIX) || it.name.startsWith(TYPES_PREFIX)) && stubPkgsFilter(it) }
+        .mapNotNull { stubPkg ->
+          (nameToPkg[stubPkg.name.removeSuffix(STUBS_SUFFIX)] ?: nameToPkg[stubPkg.name.removePrefix(TYPES_PREFIX)])
+            ?.let { it to stubPkg }
+        }
         .filter {
           val runtimePkgName = it.first.name
           val requirement = it.second.requirements.firstOrNull { req -> req.equals(runtimePkgName) } ?: return@filter false

python/src/com/jetbrains/python/codeInsight/typing/PyStubPackagesLoader.kt

@@ -49,19 +49,17 @@ private fun sourceToStubPackagesAvailableToInstall(sourceToInstalledRuntimeAndSt
                                                    availablePackages: List<RepoPackage>): Map<String, Set<RepoPackage>> {
   if (sourceToInstalledRuntimeAndStubPkgs.isEmpty()) return emptyMap()
 
-  val stubPkgsAvailableToInstall = mutableMapOf<String, RepoPackage>()
-  availablePackages.forEach { if (it.name.endsWith(STUBS_SUFFIX)) stubPkgsAvailableToInstall[it.name] = it }
+  val stubPkgsAvailableToInstall = availablePackages.asSequence()
+    .filter { it.name.endsWith(STUBS_SUFFIX) || it.name.startsWith(TYPES_PREFIX) }
+    .associateBy { it.name }
 
-  val result = mutableMapOf<String, Set<RepoPackage>>()
-  sourceToInstalledRuntimeAndStubPkgs.forEach { (source, runtimeAndStubPkgs) ->
-    result[source] = runtimeAndStubPkgs
+  return sourceToInstalledRuntimeAndStubPkgs.mapValues { (_, runtimeAndStubPkgs) ->
+    runtimeAndStubPkgs
       .asSequence()
       .filter { it.second == null }
-      .mapNotNull { stubPkgsAvailableToInstall["${it.first.name}$STUBS_SUFFIX"] }
+      .mapNotNull { stubPkgsAvailableToInstall["${it.first.name}$STUBS_SUFFIX"] ?: stubPkgsAvailableToInstall["$TYPES_PREFIX${it.first.name}"] }
       .toSet()
   }
-
-  return result
 }
 
 private fun loadRequirementsAndExtraArgs(sourceToStubPackagesAvailableToInstall: Map<String, Set<RepoPackage>>,
@@ -83,12 +81,14 @@ private fun installedRuntimeAndStubPackages(pkgName: String, installedPackages:
   var runtime: PyPackage? = null
   var stub: PyPackage? = null
   val stubPkgName = "$pkgName$STUBS_SUFFIX"
+  val typesPkgName = "$TYPES_PREFIX$pkgName"
 
   for (pkg in installedPackages) {
     val name = pkg.name
 
     if (name == pkgName) runtime = pkg
     if (name == stubPkgName) stub = pkg
+    if (name == typesPkgName) stub = pkg
   }
 
   return if (runtime == null) null else runtime to stub