Simplify stubs management without types.requirements()

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")

examples/demo/manual_stubs/explicit/BUILD.bazel

@@ -0,0 +1,13 @@
+load("@rules_python//python:py_library.bzl", "py_library")
+
+py_library(
+    name = "lib",
+    srcs = glob(["lib/**/*.py"]),
+    deps = [":foo-stubs"],
+)
+
+py_library(
+    name = "foo-stubs",
+    imports = ["."],
+    pyi_srcs = glob(["foo/**/*.pyi"]),
+)

examples/demo/manual_stubs/explicit/foo/__init__.pyi

@@ -0,0 +1,2 @@
+class Foo:
+    pass

examples/demo/manual_stubs/explicit/lib/main.py

@@ -0,0 +1 @@
+from foo import Foo

examples/demo/manual_stubs/implicit/BUILD.bazel

@@ -0,0 +1,19 @@
+load("@pip//:requirements.bzl", "requirement")
+load("@rules_python//python:py_library.bzl", "py_library")
+
+py_library(
+    name = "lib",
+    srcs = glob(["lib/**/*.py"]),
+    deps = [
+        # actual dependency here doesn't matter, we just need some dependency
+        # to implicitly provide stubs for; see `py.bzl` for the hook-up
+        requirement("six"),
+    ],
+)
+
+py_library(
+    name = "foo-stubs",
+    imports = ["."],
+    pyi_srcs = glob(["foo/**/*.pyi"]),
+    visibility = ["//visibility:public"],
+)

examples/demo/manual_stubs/implicit/foo/__init__.pyi

@@ -0,0 +1,2 @@
+class Foo:
+    pass

examples/demo/manual_stubs/implicit/lib/main.py

@@ -0,0 +1 @@
+from foo import Foo

examples/demo/py.bzl

@@ -1,6 +1,14 @@
 "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", "requirement")
+load("@rules_mypy//mypy:mypy.bzl", "load_stubs", "mypy")
 
-mypy_aspect = mypy(types = types)
+stubs = load_stubs(
+    requirements = all_requirements,
+    overrides = {
+        # See manual_stubs/implicit/
+        requirement("six"): "@@//manual_stubs/implicit:foo-stubs",
+    },
+)
+
+mypy_aspect = mypy(stubs = stubs)

examples/demo/requirements.in

@@ -1,3 +1,7 @@
-cachetools~=5.4.0
+# ----- thm/ dependencies ----- #
 numpy~=1.26.4
+cachetools~=5.4.0
 types-cachetools~=5.4.0.20240717
+
+# ---- manual_stubs/ dependencies ----- #
+six==1.17.0

examples/demo/requirements.txt

@@ -44,6 +44,10 @@ numpy==1.26.4 \
     --hash=sha256:f870204a840a60da0b12273ef34f7051e98c3b5961b61b0c2c1be6dfd64fbcd3 \
     --hash=sha256:ffa75af20b44f8dba823498024771d5ac50620e6915abac414251bd971b4529f
     # via -r requirements.in
+six==1.17.0 \
+    --hash=sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274 \
+    --hash=sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81
+    # via -r requirements.in
 types-cachetools==5.4.0.20240717 \
     --hash=sha256:1eae90c48760bac44ab89108be938e8ce1d740910f2d4b68446dcdc82763f186 \
     --hash=sha256:67c84c26df988039be68344b162afd2dd7cd3741dc08e7d67aa1954782fd2d2a

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")

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,
 )

mypy/BUILD.bazel

@@ -5,15 +5,3 @@ bzl_library(
     srcs = ["mypy.bzl"],
     visibility = ["//visibility:public"],
 )
-
-bzl_library(
-    name = "py_type_library",
-    srcs = ["py_type_library.bzl"],
-    visibility = ["//visibility:public"],
-)
-
-bzl_library(
-    name = "types",
-    srcs = ["types.bzl"],
-    visibility = ["//visibility:public"],
-)

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

mypy/private/BUILD.bazel

@@ -1,6 +1,4 @@
 load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
-load("@rules_mypy_pip//:requirements.bzl", "requirement")
-load("@rules_python//python:py_binary.bzl", "py_binary")
 load("@rules_uv//uv:pip.bzl", "pip_compile")
 load("@rules_uv//uv:venv.bzl", "create_venv")
 load(":mypy.bzl", "mypy_cli")
@@ -16,18 +14,6 @@ bzl_library(
     visibility = ["//mypy:__subpackages__"],
 )
 
-bzl_library(
-    name = "py_type_library_rules",
-    srcs = ["py_type_library.bzl"],
-    visibility = ["//mypy:__subpackages__"],
-)
-
-bzl_library(
-    name = "types_rules",
-    srcs = ["types.bzl"],
-    visibility = ["//mypy:__subpackages__"],
-)
-
 pip_compile(
     name = "generate_requirements_lock",
     requirements_in = "requirements.in",
@@ -40,14 +26,3 @@ create_venv(
 )
 
 mypy_cli(name = "mypy")
-
-py_binary(
-    name = "py_type_library",
-    srcs = ["py_type_library.py"],
-    main = "py_type_library.py",
-    python_version = "3.12",
-    visibility = ["//visibility:public"],
-    deps = [
-        requirement("click"),
-    ],
-)

mypy/private/mypy.bzl

@@ -10,7 +10,6 @@ directories.
 load("@rules_mypy_pip//:requirements.bzl", "requirement")
 load("@rules_python//python:py_binary.bzl", "py_binary")
 load("@rules_python//python:py_info.bzl", RulesPythonPyInfo = "PyInfo")
-load(":py_type_library.bzl", "PyTypeLibraryInfo")
 
 MypyCacheInfo = provider(
     doc = "Output details of the mypy build rule.",
@@ -106,12 +105,12 @@ def _mypy_impl(target, ctx):
 
     upstream_caches = []
 
-    types = []
+    stubs_deps = []
 
     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
@@ -128,20 +127,22 @@ def _mypy_impl(target, ctx):
             pyi_files.extend(dep[RulesPythonPyInfo].direct_pyi_files.to_list())
             pyi_dirs |= {"%s/%s" % (ctx.bin_dir.path, imp): None for imp in _extract_imports(dep) if imp != "site-packages" and imp != "_main"}
         depsets.append(dep.default_runfiles.files)
-        if PyTypeLibraryInfo in dep:
-            types.append(dep[PyTypeLibraryInfo].directory.path + "/site-packages")
-        elif dep.label in type_mapping:
+
+        if dep.label in type_mapping:
             continue
+
+        if dep.label.workspace_name == "":
+            for import_ in _extract_imports(dep):
+                imports_dirs[import_] = 1
+        elif dep in type_mapping.values():
+            stubs_deps.append(dep.label.workspace_root + "/site-packages")
         elif dep.label.workspace_root.startswith("external/"):
             # TODO: do we need this, still?
             external_deps[dep.label.workspace_root + "/site-packages"] = 1
             for imp in _imports(dep):
                 path = "external/{}".format(imp)
                 if path not in dep_with_stubs:
                     external_deps[path] = 1
-        elif dep.label.workspace_name == "":
-            for import_ in _extract_imports(dep):
-                imports_dirs[import_] = 1
 
         if MypyCacheInfo in dep:
             upstream_caches.append(dep[MypyCacheInfo].directory)
@@ -169,9 +170,6 @@ def _mypy_impl(target, ctx):
         # output bin_dir from `generated_dirs` first
         # https://github.com/theoremlp/rules_mypy/issues/88
         ["."] +
-        # types need to appear first in the mypy path since the module directories
-        # are the same and mypy resolves the first ones, first.
-        sorted(types) +
         sorted(external_deps) +
         sorted(imports_dirs) +
         sorted(generated_dirs) +
@@ -224,15 +222,15 @@ def _mypy_impl(target, ctx):
         outputs = outputs,
         executable = ctx.executable._mypy_cli,
         arguments = [args],
-        env = {"MYPYPATH": mypy_path} | ctx.configuration.default_shell_env | extra_env,
+        env = {"MYPYPATH": mypy_path, "PYTHONPATH": ":".join(sorted(stubs_deps))} | ctx.configuration.default_shell_env | extra_env,
     )
 
     return result_info
 
 def mypy(
         mypy_cli = None,
         mypy_ini = None,
-        types = None,
+        stubs = None,
         cache = True,
         color = True,
         suppression_tags = None,
@@ -244,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
@@ -264,8 +254,6 @@ def mypy(
     Returns:
         a mypy aspect.
     """
-    types = types or {}
-
     additional_attrs = {}
 
     return aspect(
@@ -286,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.