Fix over eager types-google-cloud-ndb suggestion

Fixes #15343

Author: hauntsaninja

mypy/build.py

@@ -57,7 +57,6 @@
     DecodeError,
     decode_python_encoding,
     get_mypy_comments,
-    get_top_two_prefixes,
     hash_digest,
     is_stub_package_file,
     is_sub_path,
@@ -97,7 +96,7 @@
     is_legacy_bundled_package,
     legacy_bundled_packages,
     non_bundled_packages,
-    stub_package_name,
+    stub_distribution_name,
 )
 from mypy.types import Type
 from mypy.typestate import reset_global_state, type_state
@@ -2664,14 +2663,15 @@ def find_module_and_diagnose(
         # search path or the module has not been installed.
 
         ignore_missing_imports = options.ignore_missing_imports
-        top_level, second_level = get_top_two_prefixes(id)
+
+        id_components = id.split(".")
         # Don't honor a global (not per-module) ignore_missing_imports
         # setting for modules that used to have bundled stubs, as
         # otherwise updating mypy can silently result in new false
         # negatives. (Unless there are stubs but they are incomplete.)
         global_ignore_missing_imports = manager.options.ignore_missing_imports
         if (
-            (is_legacy_bundled_package(top_level) or is_legacy_bundled_package(second_level))
+            any(".".join(id_components[:i]) in legacy_bundled_packages for i in range(len(id_components), 0, -1))
             and global_ignore_missing_imports
             and not options.ignore_missing_imports_per_module
             and result is ModuleNotFoundReason.APPROVED_STUBS_NOT_INSTALLED
@@ -2780,15 +2780,19 @@ def module_not_found(
         daemon = manager.options.fine_grained_incremental
         msg, notes = reason.error_message_templates(daemon)
         errors.report(line, 0, msg.format(module=target), code=codes.IMPORT)
-        top_level, second_level = get_top_two_prefixes(target)
-        if second_level in legacy_bundled_packages or second_level in non_bundled_packages:
-            top_level = second_level
+
+        components = target.split(".")
+        for i in range(len(components), 0, -1):
+            module = ".".join(components[:i])
+            if module in legacy_bundled_packages or module in non_bundled_packages:
+                break
+
         for note in notes:
             if "{stub_dist}" in note:
-                note = note.format(stub_dist=stub_package_name(top_level))
+                note = note.format(stub_dist=stub_distribution_name(module))
             errors.report(line, 0, note, severity="note", only_once=True, code=codes.IMPORT)
         if reason is ModuleNotFoundReason.APPROVED_STUBS_NOT_INSTALLED:
-            manager.missing_stub_packages.add(stub_package_name(top_level))
+            manager.missing_stub_packages.add(stub_distribution_name(module))
     errors.set_import_context(save_import_context)
 
 

mypy/modulefinder.py

@@ -337,14 +337,9 @@ def _find_module_non_stub_helper(
             # If this is not a directory then we can't traverse further into it
             if not self.fscache.isdir(dir_path):
                 break
-        if approved_stub_package_exists(components[0]):
-            if len(components) == 1 or (
-                self.find_module(components[0])
-                is ModuleNotFoundReason.APPROVED_STUBS_NOT_INSTALLED
-            ):
+        for i in range(len(components), 0, -1):
+            if approved_stub_package_exists(".".join(components[:i])):
                 return ModuleNotFoundReason.APPROVED_STUBS_NOT_INSTALLED
-        if approved_stub_package_exists(".".join(components[:2])):
-            return ModuleNotFoundReason.APPROVED_STUBS_NOT_INSTALLED
         if plausible_match:
             return ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS
         else:

mypy/stubinfo.py

@@ -9,7 +9,7 @@ def approved_stub_package_exists(prefix: str) -> bool:
     return is_legacy_bundled_package(prefix) or prefix in non_bundled_packages
 
 
-def stub_package_name(prefix: str) -> str:
+def stub_distribution_name(prefix: str) -> str:
     return legacy_bundled_packages.get(prefix) or non_bundled_packages[prefix]
 
 
@@ -116,7 +116,7 @@ def stub_package_name(prefix: str) -> str:
     "flask_sqlalchemy": "types-Flask-SQLAlchemy",
     "fpdf": "types-fpdf2",
     "gdb": "types-gdb",
-    "google.cloud": "types-google-cloud-ndb",
+    "google.cloud.ndb": "types-google-cloud-ndb",
     "hdbcli": "types-hdbcli",
     "html5lib": "types-html5lib",
     "httplib2": "types-httplib2",

mypy/util.py

@@ -308,17 +308,6 @@ def get_prefix(fullname: str) -> str:
     return fullname.rsplit(".", 1)[0]
 
 
-def get_top_two_prefixes(fullname: str) -> tuple[str, str]:
-    """Return one and two component prefixes of a fully qualified name.
-
-    Given 'a.b.c.d', return ('a', 'a.b').
-
-    If fullname has only one component, return (fullname, fullname).
-    """
-    components = fullname.split(".", 3)
-    return components[0], ".".join(components[:2])
-
-
 def correct_relative_import(
     cur_mod_id: str, relative: int, target: str, is_cur_package_init_file: bool
 ) -> tuple[str, bool]:

test-data/unit/check-modules.test

@@ -3134,17 +3134,17 @@ import google.cloud
 from google.cloud import x
 
 [case testErrorFromGoogleCloud]
-import google.cloud
+import google.cloud  # E: Cannot find implementation or library stub for module named "google.cloud" \
+                     # E: Cannot find implementation or library stub for module named "google"
 from google.cloud import x
-import google.non_existent
+import google.non_existent  # E: Cannot find implementation or library stub for module named "google.non_existent"
 from google.non_existent import x
-[out]
-main:1: error: Library stubs not installed for "google.cloud"
-main:1: note: Hint: "python3 -m pip install types-google-cloud-ndb"
-main:1: note: (or run "mypy --install-types" to install all missing stub packages)
-main:1: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
-main:1: error: Cannot find implementation or library stub for module named "google"
-main:3: error: Cannot find implementation or library stub for module named "google.non_existent"
+
+import google.cloud.ndb  # E: Library stubs not installed for "google.cloud.ndb" \
+                         # N: Hint: "python3 -m pip install types-google-cloud-ndb" \
+                         # N: (or run "mypy --install-types" to install all missing stub packages) \
+                         # N: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
+from google.cloud import ndb
 
 [case testMissingSubmoduleOfInstalledStubPackage]
 import bleach.xyz