Thanks to visit codestin.com
Credit goes to github.com

Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 19, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,8 @@ Unreleased changes template.
([#1647](https://github.com/bazelbuild/rules_python/issues/1647))
* (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
Expand Down
36 changes: 34 additions & 2 deletions python/private/py_executable.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,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(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: what about main becoming main_module if it cannot find the file in srcs?

That said, then we would have worse error messages, so this solution is also good.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd rather not because it's not clear how to transform the file name into a module name.

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(
Expand Down Expand Up @@ -638,14 +656,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,
},
Expand Down Expand Up @@ -929,7 +952,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)

Expand Down Expand Up @@ -1049,6 +1075,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")
Expand Down
114 changes: 69 additions & 45 deletions python/private/stage2_bootstrap_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 =====

Expand Down Expand Up @@ -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."""

Expand All @@ -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)
Expand Down Expand Up @@ -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)


Expand Down
9 changes: 9 additions & 0 deletions tests/bootstrap_impls/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
17 changes: 17 additions & 0 deletions tests/bootstrap_impls/main_module.py
Original file line number Diff line number Diff line change
@@ -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__}")