feat(rules): add main_module attribute to run a module name (python -m) (#2671)

rickeylev · web-flow · commit 032f6aa738a6 · 2025-03-19T06:53:06.000Z
This implements the ability to run a module name instead of a file path, aka `python -m` style of invocation. This allows a binary/test to specify what the main module is without having to have a direct dependency on the entry point file. As a side effect, the `srcs` attribute is no longer required. Fixes #2539
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -99,6 +99,8 @@ Unreleased changes template.
   Only applicable for {obj}`--bootstrap_impl=script`.
 * (rules) Added {obj}`interpreter_args` attribute to `py_binary` and `py_test`,
   which allows pass arguments to the interpreter before the regular args.
+* (rules) Added {obj}`main_module` attribute to `py_binary` and `py_test`,
+  which allows specifying a module name to run (i.e. `python -m <module>`).
 
 {#v0-0-0-removed}
 ### Removed
diff --git a/python/private/py_executable.bzl b/python/private/py_executable.bzl
@@ -130,6 +130,24 @@ Optional; the name of the source file that is the main entry point of the
 application. This file must also be listed in `srcs`. If left unspecified,
 `name`, with `.py` appended, is used instead. If `name` does not match any
 filename in `srcs`, `main` must be specified.
+
+This is mutually exclusive with {obj}`main_module`.
+""",
+        ),
+        "main_module": lambda: attrb.String(
+            doc = """
+Module name to execute as the main program.
+
+When set, `srcs` is not required, and it is assumed the module is
+provided by a dependency.
+
+See https://docs.python.org/3/using/cmdline.html#cmdoption-m for more
+information about running modules as the main program.
+
+This is mutually exclusive with {obj}`main`.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
 """,
         ),
         "pyc_collection": lambda: attrb.String(
@@ -642,14 +660,19 @@ def _create_stage2_bootstrap(
 
     template = runtime.stage2_bootstrap_template
 
+    if main_py:
+        main_py_path = "{}/{}".format(ctx.workspace_name, main_py.short_path)
+    else:
+        main_py_path = ""
     ctx.actions.expand_template(
         template = template,
         output = output,
         substitutions = {
             "%coverage_tool%": _get_coverage_tool_runfiles_path(ctx, runtime),
             "%import_all%": "True" if ctx.fragments.bazel_py.python_import_all_repositories else "False",
             "%imports%": ":".join(imports.to_list()),
-            "%main%": "{}/{}".format(ctx.workspace_name, main_py.short_path),
+            "%main%": main_py_path,
+            "%main_module%": ctx.attr.main_module,
             "%target%": str(ctx.label),
             "%workspace_name%": ctx.workspace_name,
         },
@@ -933,7 +956,10 @@ def py_executable_base_impl(ctx, *, semantics, is_test, inherited_environment =
     """
     _validate_executable(ctx)
 
-    main_py = determine_main(ctx)
+    if not ctx.attr.main_module:
+        main_py = determine_main(ctx)
+    else:
+        main_py = None
     direct_sources = filter_to_py_srcs(ctx.files.srcs)
     precompile_result = semantics.maybe_precompile(ctx, direct_sources)
 
@@ -1053,6 +1079,12 @@ def _validate_executable(ctx):
     if ctx.attr.python_version == "PY2":
         fail("It is not allowed to use Python 2")
 
+    if ctx.attr.main and ctx.attr.main_module:
+        fail((
+            "Only one of main and main_module can be set, got: " +
+            "main={}, main_module={}"
+        ).format(ctx.attr.main, ctx.attr.main_module))
+
 def _declare_executable_file(ctx):
     if target_platform_has_any_constraint(ctx, ctx.attr._windows_constraints):
         executable = ctx.actions.declare_file(ctx.label.name + ".exe")
diff --git a/python/private/stage2_bootstrap_template.py b/python/private/stage2_bootstrap_template.py
@@ -26,7 +26,11 @@
 # We just put them in one place so its easy to tell which are used.
 
 # Runfiles-relative path to the main Python source file.
-MAIN = "%main%"
+# Empty if MAIN_MODULE is used
+MAIN_PATH = "%main%"
+
+# Module name to execute. Empty if MAIN is used.
+MAIN_MODULE = "%main_module%"
 
 # ===== Template substitutions end =====
 
@@ -249,7 +253,7 @@ def unresolve_symlinks(output_filename):
         os.unlink(unfixed_file)
 
 
-def _run_py(main_filename, *, args, cwd=None):
+def _run_py_path(main_filename, *, args, cwd=None):
     # type: (str, str, list[str], dict[str, str]) -> ...
     """Executes the given Python file using the various environment settings."""
 
@@ -269,6 +273,11 @@ def _run_py(main_filename, *, args, cwd=None):
         sys.argv = orig_argv
 
 
+def _run_py_module(module_name):
+    # Match `python -m` behavior, so modify sys.argv and the run name
+    runpy.run_module(module_name, alter_sys=True, run_name="__main__")
+
+
 @contextlib.contextmanager
 def _maybe_collect_coverage(enable):
     print_verbose_coverage("enabled:", enable)
@@ -356,64 +365,79 @@ def main():
     print_verbose("initial environ:", mapping=os.environ)
     print_verbose("initial sys.path:", values=sys.path)
 
-    main_rel_path = MAIN
-    if is_windows():
-        main_rel_path = main_rel_path.replace("/", os.sep)
-
-    module_space = find_runfiles_root(main_rel_path)
-    print_verbose("runfiles root:", module_space)
-
-    # Recreate the "add main's dir to sys.path[0]" behavior to match the
-    # system-python bootstrap / typical Python behavior.
-    #
-    # Without safe path enabled, when `python foo/bar.py` is run, python will
-    # resolve the foo/bar.py symlink to its real path, then add the directory
-    # of that path to sys.path. But, the resolved directory for the symlink
-    # depends on if the file is generated or not.
-    #
-    # When foo/bar.py is a source file, then it's a symlink pointing
-    # back to the client source directory. This means anything from that source
-    # directory becomes importable, i.e. most code is importable.
-    #
-    # When foo/bar.py is a generated file, then it's a symlink pointing to
-    # somewhere under bazel-out/.../bin, i.e. where generated files are. This
-    # means only other generated files are importable (not source files).
-    #
-    # To replicate this behavior, we add main's directory within the runfiles
-    # when safe path isn't enabled.
-    if not getattr(sys.flags, "safe_path", False):
-        prepend_path_entries = [
-            os.path.join(module_space, os.path.dirname(main_rel_path))
-        ]
+    main_rel_path = None
+    # todo: things happen to work because find_runfiles_root
+    # ends up using stage2_bootstrap, and ends up computing the proper
+    # runfiles root
+    if MAIN_PATH:
+        main_rel_path = MAIN_PATH
+        if is_windows():
+            main_rel_path = main_rel_path.replace("/", os.sep)
+
+        runfiles_root = find_runfiles_root(main_rel_path)
     else:
-        prepend_path_entries = []
+        runfiles_root = find_runfiles_root("")
+
+    print_verbose("runfiles root:", runfiles_root)
 
-    runfiles_envkey, runfiles_envvalue = runfiles_envvar(module_space)
+    runfiles_envkey, runfiles_envvalue = runfiles_envvar(runfiles_root)
     if runfiles_envkey:
         os.environ[runfiles_envkey] = runfiles_envvalue
 
-    main_filename = os.path.join(module_space, main_rel_path)
-    main_filename = get_windows_path_with_unc_prefix(main_filename)
-    assert os.path.exists(main_filename), (
-        "Cannot exec() %r: file not found." % main_filename
-    )
-    assert os.access(main_filename, os.R_OK), (
-        "Cannot exec() %r: file not readable." % main_filename
-    )
+    if MAIN_PATH:
+        # Recreate the "add main's dir to sys.path[0]" behavior to match the
+        # system-python bootstrap / typical Python behavior.
+        #
+        # Without safe path enabled, when `python foo/bar.py` is run, python will
+        # resolve the foo/bar.py symlink to its real path, then add the directory
+        # of that path to sys.path. But, the resolved directory for the symlink
+        # depends on if the file is generated or not.
+        #
+        # When foo/bar.py is a source file, then it's a symlink pointing
+        # back to the client source directory. This means anything from that source
+        # directory becomes importable, i.e. most code is importable.
+        #
+        # When foo/bar.py is a generated file, then it's a symlink pointing to
+        # somewhere under bazel-out/.../bin, i.e. where generated files are. This
+        # means only other generated files are importable (not source files).
+        #
+        # To replicate this behavior, we add main's directory within the runfiles
+        # when safe path isn't enabled.
+        if not getattr(sys.flags, "safe_path", False):
+            prepend_path_entries = [
+                os.path.join(runfiles_root, os.path.dirname(main_rel_path))
+            ]
+        else:
+            prepend_path_entries = []
+
+        main_filename = os.path.join(runfiles_root, main_rel_path)
+        main_filename = get_windows_path_with_unc_prefix(main_filename)
+        assert os.path.exists(main_filename), (
+            "Cannot exec() %r: file not found." % main_filename
+        )
+        assert os.access(main_filename, os.R_OK), (
+            "Cannot exec() %r: file not readable." % main_filename
+        )
 
-    sys.stdout.flush()
+        sys.stdout.flush()
 
-    sys.path[0:0] = prepend_path_entries
+        sys.path[0:0] = prepend_path_entries
+    else:
+        main_filename = None
 
     if os.environ.get("COVERAGE_DIR"):
         import _bazel_site_init
+
         coverage_enabled = _bazel_site_init.COVERAGE_SETUP
     else:
         coverage_enabled = False
 
     with _maybe_collect_coverage(enable=coverage_enabled):
-        # The first arg is this bootstrap, so drop that for the re-invocation.
-        _run_py(main_filename, args=sys.argv[1:])
+        if MAIN_PATH:
+            # The first arg is this bootstrap, so drop that for the re-invocation.
+            _run_py_path(main_filename, args=sys.argv[1:])
+        else:
+            _run_py_module(MAIN_MODULE)
         sys.exit(0)
 
 
diff --git a/tests/bootstrap_impls/BUILD.bazel b/tests/bootstrap_impls/BUILD.bazel
@@ -107,6 +107,15 @@ py_reconfig_test(
     main = "sys_path_order_test.py",
 )
 
+py_reconfig_test(
+    name = "main_module_test",
+    srcs = ["main_module.py"],
+    bootstrap_impl = "script",
+    imports = ["."],
+    main_module = "tests.bootstrap_impls.main_module",
+    target_compatible_with = SUPPORTS_BOOTSTRAP_SCRIPT,
+)
+
 sh_py_run_test(
     name = "inherit_pythonsafepath_env_test",
     bootstrap_impl = "script",
diff --git a/tests/bootstrap_impls/main_module.py b/tests/bootstrap_impls/main_module.py
@@ -0,0 +1,17 @@
+import sys
+import unittest
+
+
+class MainModuleTest(unittest.TestCase):
+    def test_run_as_module(self):
+        self.assertIsNotNone(__spec__, "__spec__ was none")
+        # If not run as a module, __spec__ is None
+        self.assertNotEqual(__name__, __spec__.name)
+        self.assertEqual(__spec__.name, "tests.bootstrap_impls.main_module")
+
+
+if __name__ == "__main__":
+    unittest.main()
+else:
+    # Guard against running it as a module in a non-main way.
+    sys.exit(f"__name__ should be __main__, got {__name__}")
