feat(gazelle): Include types/stubs packages that have them

Author: ewianda

gazelle/modules_mapping/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@rules_python//python:defs.bzl", "py_binary")
+load("@rules_python//python:defs.bzl", "py_binary", "py_test")
 
 # gazelle:exclude *.py
 
@@ -8,6 +8,15 @@ py_binary(
     visibility = ["//visibility:public"],
 )
 
+py_test(
+    name = "test_generator",
+    srcs = ["test_generator.py"],
+    data = glob(["testdata/**"]),
+    imports = ["."],
+    main = "test_generator.py",
+    deps = [":generator"],
+)
+
 filegroup(
     name = "distribution",
     srcs = glob(["**"]),

gazelle/modules_mapping/generator.py

@@ -35,6 +35,11 @@ def __init__(self, stderr, output_file, excluded_patterns):
     # dig_wheel analyses the wheel .whl file determining the modules it provides
     # by looking at the directory structure.
     def dig_wheel(self, whl):
+        # Skip stubs and types wheels.
+        wheel_name = get_wheel_name(whl)
+        if wheel_name.endswith(("_stubs", "_types")):
+            self.mapping[wheel_name.lower()] = wheel_name.lower()
+            return
         with zipfile.ZipFile(whl, "r") as zip_file:
             for path in zip_file.namelist():
                 if is_metadata(path):

gazelle/modules_mapping/test_generator.py

@@ -0,0 +1,41 @@
+import pathlib
+import unittest
+
+from generator import Generator
+
+
+class GeneratorTest(unittest.TestCase):
+    def test_generator(self):
+        whl = pathlib.Path(
+            pathlib.Path(__file__).parent, "testdata", "pytest-7.1.1-py3-none-any.whl"
+        )
+        gen = Generator(None, None, {})
+        gen.dig_wheel(whl)
+        self.assertLessEqual(
+            {
+                "_pytest": "pytest",
+                "_pytest.__init__": "pytest",
+                "_pytest._argcomplete": "pytest",
+                "_pytest.config.argparsing": "pytest",
+            }.items(),
+            gen.mapping.items(),
+        )
+
+    def test_stub_generator(self):
+        whl = pathlib.Path(
+            pathlib.Path(__file__).parent,
+            "testdata",
+            "django_types-0.15.0-py3-none-any.whl",
+        )
+        gen = Generator(None, None, {})
+        gen.dig_wheel(whl)
+        self.assertLessEqual(
+            {
+                "django_types": "django_types",
+            }.items(),
+            gen.mapping.items(),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

gazelle/modules_mapping/testdata/django_types-0.15.0-py3-none-any.whl

@@ -189,8 +189,18 @@ func (py *Resolver) Resolve(
 						continue MODULES_LOOP
 					}
 				} else {
-					if dep, ok := cfg.FindThirdPartyDependency(moduleName); ok {
+					if dep, distributionName, ok := cfg.FindThirdPartyDependency(moduleName); ok {
 						deps.Add(dep)
+						// Add the type and stub dependencies if they exist.
+						typeModule := fmt.Sprintf("%s_types", strings.ToLower(distributionName))
+						if dep, _, ok := cfg.FindThirdPartyDependency(typeModule); ok {
+							deps.Add(dep)
+
+						}
+						stubModule := fmt.Sprintf("%s_stubs", strings.ToLower(distributionName))
+						if dep, _, ok := cfg.FindThirdPartyDependency(stubModule); ok {
+							deps.Add(dep)
+						}
 						if explainDependency == dep {
 							log.Printf("Explaining dependency (%s): "+
 								"in the target %q, the file %q imports %q at line %d, "+

gazelle/modules_mapping/testdata/pytest-7.1.1-py3-none-any.whl

@@ -278,7 +278,7 @@ func (c *Config) SetGazelleManifest(gazelleManifest *manifest.Manifest) {
 // FindThirdPartyDependency scans the gazelle manifests for the current config
 // and the parent configs up to the root finding if it can resolve the module
 // name.
-func (c *Config) FindThirdPartyDependency(modName string) (string, bool) {
+func (c *Config) FindThirdPartyDependency(modName string) (string, string, bool) {
 	for currentCfg := c; currentCfg != nil; currentCfg = currentCfg.parent {
 		if currentCfg.gazelleManifest != nil {
 			gazelleManifest := currentCfg.gazelleManifest
@@ -291,11 +291,11 @@ func (c *Config) FindThirdPartyDependency(modName string) (string, bool) {
 				}
 
 				lbl := currentCfg.FormatThirdPartyDependency(distributionRepositoryName, distributionName)
-				return lbl.String(), true
+				return lbl.String(), distributionName, true
 			}
 		}
 	}
-	return "", false
+	return "", "", false
 }
 
 // AddIgnoreFile adds a file to the list of ignored files for a given package.