Add unit test

Author: ewianda

gazelle/modules_mapping/def.bzl

@@ -32,6 +32,7 @@ def _modules_mapping_impl(ctx):
     )
     args.add("--output_file", modules_mapping.path)
     args.add_all("--exclude_patterns", ctx.attr.exclude_patterns)
+    args.add_all("--include_stub_packages", ctx.attr.include_stub_packages)
     args.add_all("--wheels", [whl.path for whl in all_wheels.to_list()])
     ctx.actions.run(
         inputs = all_wheels.to_list(),
@@ -55,6 +56,11 @@ modules_mapping = rule(
             doc = "The name for the output JSON file.",
             mandatory = False,
         ),
+        "include_stub_packages": attr.bool(
+            default = False,
+            doc = "Whether to include stub packages in the mapping.",
+            mandatory = False,
+        ),
         "wheels": attr.label_list(
             allow_files = True,
             doc = "The list of wheels, usually the 'all_whl_requirements' from @<pip_repository>//:requirements.bzl",

gazelle/modules_mapping/generator.py

@@ -25,19 +25,20 @@ class Generator:
     stderr = None
     output_file = None
     excluded_patterns = None
-    mapping = {}
 
-    def __init__(self, stderr, output_file, excluded_patterns):
+    def __init__(self, stderr, output_file, excluded_patterns, include_stub_packages):
         self.stderr = stderr
         self.output_file = output_file
         self.excluded_patterns = [re.compile(pattern) for pattern in excluded_patterns]
+        self.include_stub_packages = include_stub_packages
+        self.mapping = {}
 
     # 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")):
+        if wheel_name.endswith(("_stubs", "_types")) and self.include_stub_packages:
             self.mapping[wheel_name.lower()] = wheel_name.lower()
             return
         with zipfile.ZipFile(whl, "r") as zip_file:
@@ -150,8 +151,11 @@ def data_has_purelib_or_platlib(path):
         description="Generates the modules mapping used by the Gazelle manifest.",
     )
     parser.add_argument("--output_file", type=str)
+    parser.add_argument("--include_stub_packages", type=bool)
     parser.add_argument("--exclude_patterns", nargs="+", default=[])
     parser.add_argument("--wheels", nargs="+", default=[])
     args = parser.parse_args()
-    generator = Generator(sys.stderr, args.output_file, args.exclude_patterns)
+    generator = Generator(
+        sys.stderr, args.output_file, args.exclude_patterns, args.include_stub_packages
+    )
     exit(generator.run(args.wheels))

gazelle/modules_mapping/test_generator.py

@@ -9,7 +9,7 @@ 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 = Generator(None, None, {}, False)
         gen.dig_wheel(whl)
         self.assertLessEqual(
             {
@@ -27,7 +27,7 @@ def test_stub_generator(self):
             "testdata",
             "django_types-0.15.0-py3-none-any.whl",
         )
-        gen = Generator(None, None, {})
+        gen = Generator(None, None, {}, True)
         gen.dig_wheel(whl)
         self.assertLessEqual(
             {
@@ -36,6 +36,19 @@ def test_stub_generator(self):
             gen.mapping.items(),
         )
 
+    def test_stub_excluded(self):
+        whl = pathlib.Path(
+            pathlib.Path(__file__).parent
+            / "testdata"
+            / "django_types-0.15.0-py3-none-any.whl"
+        )
+        gen = Generator(None, None, {}, False)
+        gen.dig_wheel(whl)
+        self.assertEqual(
+            {}.items(),
+            gen.mapping.items(),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

gazelle/python/resolve.go

@@ -195,7 +195,6 @@ func (py *Resolver) Resolve(
 						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 {

gazelle/python/testdata/add_type_stub_packages/BUILD.in

@@ -0,0 +1,14 @@
+load("@rules_python//python:defs.bzl", "py_binary")
+
+py_binary(
+    name = "add_type_stub_packages_bin",
+    srcs = ["__main__.py"],
+    main = "__main__.py",
+    visibility = ["//:__subpackages__"],
+    deps = [
+        "@gazelle_python_test//boto3",
+        "@gazelle_python_test//boto3_stubs",
+        "@gazelle_python_test//django",
+        "@gazelle_python_test//django_types",
+    ],
+)

gazelle/python/testdata/add_type_stub_packages/BUILD.out

@@ -0,0 +1,4 @@
+# Add stubs to `deps` of `py_library` target
+
+This test case asserts that 
+* if a package has the corresponding stub available, it is added to the `deps` of the `py_library` target.

gazelle/python/testdata/add_type_stub_packages/README.md

@@ -0,0 +1 @@
+# This is a Bazel workspace for the Gazelle test data.

gazelle/python/testdata/add_type_stub_packages/WORKSPACE

@@ -0,0 +1,16 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import boto3
+import django

gazelle/python/testdata/add_type_stub_packages/__main__.py

@@ -0,0 +1,22 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+manifest:
+  modules_mapping:
+    boto3: boto3
+    boto3_stubs: boto3_stubs
+    django_types: django_types
+    django: Django
+
+  pip_deps_repository_name: gazelle_python_test