Replace types.requirements() with load_stubs()

brandonchinn178 · brandonchinn178 · commit 6aa90ce4a25b · 2025-08-06T10:16:15.000-07:00
diff --git a/examples/demo/MODULE.bazel b/examples/demo/MODULE.bazel
@@ -37,11 +37,3 @@ pip.parse(
     },
 )
 use_repo(pip, "pip")
-
-types = use_extension("@rules_mypy//mypy:types.bzl", "types")
-types.requirements(
-    name = "pip_types",
-    pip_requirements = "@pip//:requirements.bzl",
-    requirements_txt = "//:requirements.txt",
-)
-use_repo(types, "pip_types")
diff --git a/examples/demo/py.bzl b/examples/demo/py.bzl
@@ -1,6 +1,8 @@
 "Custom py_library rule that also runs mypy."
 
-load("@pip_types//:types.bzl", "types")
-load("@rules_mypy//mypy:mypy.bzl", "mypy")
+load("@pip//:requirements.bzl", "all_requirements")
+load("@rules_mypy//mypy:mypy.bzl", "load_stubs", "mypy")
 
-mypy_aspect = mypy(types = types)
+stubs = load_stubs(requirements = all_requirements)
+
+mypy_aspect = mypy(stubs = stubs)
diff --git a/examples/opt-in/MODULE.bazel b/examples/opt-in/MODULE.bazel
@@ -37,11 +37,3 @@ pip.parse(
     },
 )
 use_repo(pip, "pip")
-
-types = use_extension("@rules_mypy//mypy:types.bzl", "types")
-types.requirements(
-    name = "pip_types",
-    pip_requirements = "@pip//:requirements.bzl",
-    requirements_txt = "//:requirements.txt",
-)
-use_repo(types, "pip_types")
diff --git a/examples/opt-in/py.bzl b/examples/opt-in/py.bzl
@@ -1,10 +1,12 @@
 "Custom py_library rule that also runs mypy."
 
-load("@pip_types//:types.bzl", "types")
-load("@rules_mypy//mypy:mypy.bzl", "mypy")
+load("@pip//:requirements.bzl", "all_requirements")
+load("@rules_mypy//mypy:mypy.bzl", "load_stubs", "mypy")
+
+stubs = load_stubs(requirements = all_requirements)
 
 mypy_aspect = mypy(
     # only run mypy on targets with the typecheck tag
     opt_in_tags = ["typecheck"],
-    types = types,
+    stubs = stubs,
 )
diff --git a/mypy/BUILD.bazel b/mypy/BUILD.bazel
@@ -5,9 +5,3 @@ bzl_library(
     srcs = ["mypy.bzl"],
     visibility = ["//visibility:public"],
 )
-
-bzl_library(
-    name = "types",
-    srcs = ["types.bzl"],
-    visibility = ["//visibility:public"],
-)
diff --git a/mypy/mypy.bzl b/mypy/mypy.bzl
@@ -1,6 +1,13 @@
 "Public API for interacting with the mypy rule."
 
-load("//mypy/private:mypy.bzl", _mypy = "mypy", _mypy_cli = "mypy_cli")
+load(
+    "//mypy/private:mypy.bzl",
+    _load_stubs = "load_stubs",
+    _mypy = "mypy",
+    _mypy_cli = "mypy_cli",
+)
+
+load_stubs = _load_stubs
 
 # re-export mypy aspect factory
 mypy = _mypy
diff --git a/mypy/private/BUILD.bazel b/mypy/private/BUILD.bazel
@@ -14,12 +14,6 @@ bzl_library(
     visibility = ["//mypy:__subpackages__"],
 )
 
-bzl_library(
-    name = "types_rules",
-    srcs = ["types.bzl"],
-    visibility = ["//mypy:__subpackages__"],
-)
-
 pip_compile(
     name = "generate_requirements_lock",
     requirements_in = "requirements.in",
diff --git a/mypy/private/mypy.bzl b/mypy/private/mypy.bzl
@@ -109,8 +109,8 @@ def _mypy_impl(target, ctx):
 
     depsets = []
 
-    type_mapping = dict(zip([k.label for k in ctx.attr._types_keys], ctx.attr._types_values))
-    dep_with_stubs = [_.label.workspace_root + "/site-packages" for _ in ctx.attr._types_keys]
+    type_mapping = dict(zip([k.label for k in ctx.attr._stubs_keys], ctx.attr._stubs_values))
+    dep_with_stubs = [_.label.workspace_root + "/site-packages" for _ in ctx.attr._stubs_keys]
     additional_types = [
         type_mapping[dep.label]
         for dep in ctx.rule.attr.deps
@@ -230,7 +230,7 @@ def _mypy_impl(target, ctx):
 def mypy(
         mypy_cli = None,
         mypy_ini = None,
-        types = None,
+        stubs = None,
         cache = True,
         color = True,
         suppression_tags = None,
@@ -242,15 +242,7 @@ def mypy(
         mypy_cli:   (optional) a replacement mypy_cli to use (recommended to produce
                     with mypy_cli macro)
         mypy_ini:   (optional) mypy_ini file to use
-        types:      (optional) a dict of dependency label to types dependency label
-                    example:
-                    ```
-                    {
-                        requirement("cachetools"): requirement("types-cachetools"),
-                    }
-                    ```
-                    Use the types extension to create this map for a requirements.in
-                    or requirements.txt file.
+        stubs:      (optional) result from load_stubs()
         cache:      (optional, default True) propagate the mypy cache
         color:      (optional, default True) use color in mypy output
         suppression_tags: (optional, default ["no-mypy"]) tags that suppress running
@@ -262,8 +254,6 @@ def mypy(
     Returns:
         a mypy aspect.
     """
-    types = types or {}
-
     additional_attrs = {}
 
     return aspect(
@@ -284,15 +274,61 @@ def mypy(
             ),
             # pass the dict[Label, Label] in parts because Bazel doesn't have
             # this kind of attr to pass naturally
-            "_types_keys": attr.label_list(default = types.keys()),
-            "_types_values": attr.label_list(default = types.values()),
+            "_stubs_keys": attr.label_list(default = stubs.mapping.keys() if stubs else []),
+            "_stubs_values": attr.label_list(default = stubs.mapping.values() if stubs else []),
             "_suppression_tags": attr.string_list(default = suppression_tags or ["no-mypy"]),
             "_opt_in_tags": attr.string_list(default = opt_in_tags or []),
             "cache": attr.bool(default = cache),
             "color": attr.bool(default = color),
         } | additional_attrs,
     )
 
+def _load_stubs_from_requirements(requirements):
+    # for a package "foo-bar", maps "foo_bar" to "@pip//foo_bar:pkg"
+    parsed_reqs = {}
+    for req in requirements:
+        parsed_reqs[Label(req).package] = req
+
+    stubs = {}
+    for req in requirements:
+        req_name = Label(req).package
+
+        if req_name.endswith("_stubs"):
+            base_req_name = req_name.removesuffix("_stubs")
+        elif req_name.startswith("types_"):
+            base_req_name = req_name.removeprefix("types_")
+        else:
+            continue
+
+        if base_req_name in parsed_reqs:
+            base_req = parsed_reqs[base_req_name]
+            stubs[base_req] = req
+
+    return stubs
+
+def load_stubs(requirements = [], overrides = {}):
+    """
+    Generate a mapping of labels to their stubs label. For example:
+
+    ```
+    {
+        requirement("cachetools"): requirement("types-cachetools"),
+    }
+    ```
+
+    This can be detected automatically via `requirements`, manually via `overrides`,
+    or a combination of both.
+
+    Args:
+        requirements: (optional) the full list of requirements from "@pip//:requirements.bzl"
+                      to automatically detect stubs from.
+        overrides:    (optional) explicitly specified stubs mapping
+    """
+    stubs = _load_stubs_from_requirements(requirements)
+    return struct(
+        mapping = stubs | overrides,
+    )
+
 def mypy_cli(name, deps = None, mypy_requirement = None, python_version = "3.12", tags = None):
     """
     Produce a custom mypy executable for use with the mypy build rule.
diff --git a/mypy/private/types.bzl b/mypy/private/types.bzl
diff --git a/mypy/types.bzl b/mypy/types.bzl
diff --git a/readme.md b/readme.md
@@ -24,32 +24,19 @@ Setup is significantly easier with bzlmod, we recommend and predominantly suppor
 bazel_dep(name = "rules_mypy", version = "0.0.0")
 ```
 
-**Optionally, configure a types repository:**
-
-Many Python packages have separately published types/stubs packages. While mypy (and these rules) will work without including these types, this ruleset provides some utilities for leveraging these types to improve mypy's type checking.
-
-```starlark
-types = use_extension("@rules_mypy//mypy:types.bzl", "types")
-types.requirements(
-    name = "pip_types",
-    # `@pip` in the next line corresponds to the `hub_name` when using
-    # rules_python's `pip.parse(...)`.
-    pip_requirements = "@pip//:requirements.bzl",
-    # also legal to pass a `requirements.in` here
-    requirements_txt = "//:requirements.txt",
-)
-use_repo(types, "pip_types")
-```
-
 **Configure `mypy_aspect`.**
 
 Define a new aspect in a `.bzl` file (such as `./tools/aspects.bzl`):
 
 ```starlark
-load("@pip_types//:types.bzl", "types")
-load("@rules_mypy//mypy:mypy.bzl", "mypy")
+load("@pip//:requirements.bzl", "all_requirements", "requirement")
+load("@rules_mypy//mypy:mypy.bzl", "load_stubs", "mypy")
+
+stubs = load_stubs(
+    # See "Specifying stubs"
+)
 
-mypy_aspect = mypy(types = types)
+mypy_aspect = mypy(stubs = stubs)
 ```
 
 Update your `.bazelrc` to include this new aspect:
@@ -62,6 +49,64 @@ build --aspects //tools:aspects.bzl%mypy_aspect
 build --output_groups=+mypy
 ```
 
+### Specifying stubs
+
+If a third-party library does not contain type hints, it likely has type stubs defined in a `types-somelib` or `somelib-stubs` library. If it does not, you will need to create your own [stub files](https://mypy.readthedocs.io/en/stable/stubs.html) for the dependency.
+
+Whether the stub files are provided by a `types-somelib` library or written yourself, Bazel needs to be told that typechecking your code depends on `types-somelib`. This can be done in one of two ways:
+
+1. Configure `rules_mypy` to automatically add `types-somelib` to the typechecking environment when it sees `somelib` as a dependency
+    * This is the recommended approach, and should be sufficient in most cases
+    * If any stubs libraries has dependencies (e.g. `types-boto3`), this approach will not work
+
+1. Add both `requirement("somelib")` and `requirement("types-somelib")` as dependencies
+
+The rest of this section will focus on the first approach, which is done via the `stubs` parameter to the `mypy()` aspect, which should be constructed with the `load_stubs()` helper.
+
+In the common case, a library has type hints provided at `types-<name>` or `<name>-stubs`; the `requirements` parameter to `load_stubs()` automatically detects these libraries:
+
+```starlark
+# load the requirements.bzl file from the repo you configured with pip.parse()
+load("@pip//:requirements.bzl", "all_requirements")
+
+stubs = load_stubs(requirements = all_requirements)
+```
+
+Some libraries have a different stubs library name, e.g. `grpc-stubs` is the stubs library for `grpcio`. These cases need to be manually specified:
+
+```starlark
+load("@pip//:requirements.bzl", "requirement")
+
+stubs = load_stubs(
+    overrides = {
+        requirement("grpcio"): requirement("grpc-stubs"),
+    },
+)
+```
+
+If you need to write your own stubs, you can define a new `py_library` target and specify it in `load_stubs`:
+
+```starlark
+# stubs/BUILD.bazel
+
+load("@rules_python//python:py_library.bzl", "py_library")
+
+py_library(
+    name = "kafka-python",
+    imports = ["."],
+    pyi_srcs = ["kafka/__init__.pyi"],
+    visibility = ["//visibility:public"],
+)
+
+# aspects.bzl
+
+stubs = load_stubs(
+    overrides = {
+        requirement("kafka-python"): "@@//stubs:kafka-python",
+    },
+)
+```
+
 ## Customizing mypy
 
 ### Configuring mypy with mypy.ini
