fix: handle py_binary#main like rules_python does

py/BUILD.bazel

@@ -20,8 +20,10 @@ bzl_library(
     visibility = ["//visibility:public"],
     deps = [
         "//py/private:py_binary",
+        "//py/private:py_executable",
         "//py/private:py_library",
         "//py/private:py_pytest_main",
         "//py/private:py_wheel",
+        "@aspect_bazel_lib//lib:utils",
     ],
 )

py/defs.bzl

@@ -1,6 +1,8 @@
 "Public API re-exports"
 
+load("@aspect_bazel_lib//lib:utils.bzl", "propagate_common_rule_attributes")
 load("//py/private:py_binary.bzl", _py_binary = "py_binary", _py_test = "py_test")
+load("//py/private:py_executable.bzl", "determine_main")
 load("//py/private:py_library.bzl", _py_library = "py_library")
 load("//py/private:py_pytest_main.bzl", _py_pytest_main = "py_pytest_main")
 load("//py/private:py_wheel.bzl", "py_wheel_lib")
@@ -16,12 +18,20 @@ _a_struct_type = type(struct())
 _a_string_type = type("")
 
 def _py_binary_or_test(name, rule, srcs, main, imports, deps = [], resolutions = {}, **kwargs):
-    if not main and not len(srcs):
-        fail("When 'main' is not specified, 'srcs' must be non-empty")
+    # Compatibility with rules_python, see docs in py_executable.bzl
+    main_target = "_{}.find_main".format(name)
+    determine_main(
+        name = main_target,
+        target_name = name,
+        main = main,
+        srcs = srcs,
+        **propagate_common_rule_attributes(kwargs)
+    )
+
     rule(
         name = name,
         srcs = srcs,
-        main = main if main != None else srcs[0],
+        main = main_target,
         imports = imports,
         deps = deps,
         resolutions = resolutions,
@@ -46,7 +56,10 @@ def py_binary(name, srcs = [], main = None, imports = ["."], resolutions = {}, *
     Args:
         name: name of the rule
         srcs: python source files
-        main: the entry point. If absent, then the first entry in srcs is used.
+        main: the entry point.
+            Like rules_python, this is treated as a suffix of a file that should appear among the srcs.
+            If absent, then "[name].py" is tried. As a final fallback, if the srcs has a single file,
+            that is used as the main.
         imports: List of import paths to add for this binary.
         resolutions: FIXME
         **kwargs: additional named parameters to the py_binary_rule
@@ -78,6 +91,9 @@ def py_binary(name, srcs = [], main = None, imports = ["."], resolutions = {}, *
 
 def py_test(name, main = None, srcs = [], imports = ["."], **kwargs):
     "Identical to py_binary, but produces a target that can be used with `bazel test`."
+
+    # Ensure that any other targets we write will be testonly like the py_test target
+    kwargs["testonly"] = True
     _py_binary_or_test(name = name, rule = _py_test, srcs = srcs, main = main, imports = imports, **kwargs)
 
 py_wheel = rule(

py/private/BUILD.bazel

@@ -64,6 +64,12 @@ bzl_library(
     visibility = ["//py:__subpackages__"],
 )
 
+bzl_library(
+    name = "py_executable",
+    srcs = ["py_executable.bzl"],
+    visibility = ["//py:__subpackages__"],
+)
+
 bzl_library(
     name = "utils",
     srcs = ["utils.bzl"],

py/private/py_executable.bzl

@@ -0,0 +1,93 @@
+"""Some vendored helper code, copied from Bazel:
+https://github.com/bazelbuild/bazel/blob/37983b2f89cca7cc8b0bd596f4558fa36c8fbff2/src/main/starlark/builtins_bzl/common/python/py_executable.bzl
+
+There's not a way for us to reference that at runtime (via @bazel_tools, say) and even if there was
+we don't want to have to detect the version of Bazel that the user is running to know whether code
+we want to re-use will be present.
+"""
+
+def csv(values):
+    """Convert a list of strings to comma separated value string."""
+    return ", ".join(sorted(values))
+
+def _path_endswith(path, endswith):
+    # Use slash to anchor each path to prevent e.g.
+    # "ab/c.py".endswith("b/c.py") from incorrectly matching.
+    return ("/" + path).endswith("/" + endswith)
+
+def _determine_main(ctx):
+    """Determine the main entry point .py source file.
+
+    Args:
+        ctx: The rule ctx.
+
+    Returns:
+        Artifact; the main file. If one can't be found, an error is raised.
+    """
+    if ctx.attr.main:
+        # Deviation from rules_python: allow a leading colon, e.g. `main = ":my_target"`
+        proposed_main = ctx.attr.main.removeprefix(":")
+        if not proposed_main.endswith(".py"):
+            fail("main {} must end in '.py'".format(proposed_main))
+    else:
+        if ctx.attr.target_name.endswith(".py"):
+            fail("name {} must not end in '.py'".format(ctx.attr.target_name))
+        proposed_main = ctx.attr.target_name + ".py"
+
+    main_files = [src for src in ctx.files.srcs if _path_endswith(src.short_path, proposed_main)]
+
+    # Deviation from logic in rules_python: rules_py is a bit more permissive.
+    # Allow a srcs of length one to determine the main, if the target name didn't match anything.
+    if not main_files and len(ctx.files.srcs) == 1:
+        main_files = ctx.files.srcs
+
+    if not main_files:
+        if ctx.attr.main:
+            fail("could not find '{}' as specified by 'main' attribute".format(proposed_main))
+        else:
+            fail(("corresponding default '{}' does not appear in srcs. Add " +
+                  "it or override default file name with a 'main' attribute").format(
+                proposed_main,
+            ))
+
+    elif len(main_files) > 1:
+        if ctx.attr.main:
+            fail(("file name '{}' specified by 'main' attributes matches multiple files. " +
+                  "Matches: {}").format(
+                proposed_main,
+                csv([f.short_path for f in main_files]),
+            ))
+        else:
+            fail(("default main file '{}' matches multiple files in srcs. Perhaps specify " +
+                  "an explicit file with 'main' attribute? Matches were: {}").format(
+                proposed_main,
+                csv([f.short_path for f in main_files]),
+            ))
+    return main_files[0]
+
+# Adapts the function above, which we copied from rules_python, to be a standalone rule so we can
+# use it from a macro.
+# This is because we want our underlying py_binary rule to be simple: 'main' is a mandatory label.
+# The default output of this rule will be used as that label.
+def _determine_main_impl(ctx):
+    return DefaultInfo(files = depset([_determine_main(ctx)]))
+
+determine_main = rule(
+    doc = """rules_python compatibility shim: find a main file with the given name among the srcs.
+
+    From rules_python:
+    https://github.com/bazelbuild/rules_python/blob/4fe0db3cdcc063d5bdeab756e948640f3f16ae33/python/private/common/py_executable.bzl#L73
+    # TODO(b/203567235): In the Java impl, any file is allowed. While marked
+    # label, it is more treated as a string, and doesn't have to refer to
+    # anything that exists because it gets treated as suffix-search string
+    # over `srcs`.
+
+    We think these are lame semantics, however we want rules_py to be a drop-in replacement for rules_python.
+    """,
+    implementation = _determine_main_impl,
+    attrs = {
+        "target_name": attr.string(mandatory = True, doc = "The name of the py_binary or py_test we are finding a main for"),
+        "main": attr.string(doc = "Hint the user supplied as the main"),
+        "srcs": attr.label_list(allow_files = True),
+    },
+)

py/tests/external-deps/custom-macro/macro.bzl

@@ -1,13 +1,15 @@
 # buildifier: disable=module-docstring
-load("//py:defs.bzl", "py_binary", "py_library")
+load("//py:defs.bzl", "py_binary_rule", "py_library")
 
 def click_cli_binary(name, deps = [], **kwargs):
     py_library(
         name = name + "_lib",
         srcs = ["//py/tests/external-deps/custom-macro:__main__.py"],
     )
 
-    py_binary(
+    # NB: we don't use the py_binary macro here, because we want our `main` attribute to be used
+    # exactly as specified here, rather than follow rules_python semantics.
+    py_binary_rule(
         name = name,
         main = "//py/tests/external-deps/custom-macro:__main__.py",
         deps = deps + [

py/tests/py-binary/BUILD.bazel

@@ -0,0 +1,50 @@
+load("@aspect_bazel_lib//lib:testing.bzl", "assert_contains")
+load("@aspect_bazel_lib//lib:run_binary.bzl", "run_binary")
+load("//py:defs.bzl", "py_binary")
+
+#################
+# Case 1: main can be a string referencing some source file
+# Inline the syntax-sugar produced by the entry_point() helper in rules_python,
+# declaring our own binary that directly invokes a third-party entry point.
+# We happen to use pytest here only for convenience; it's a program already in our requirements.in.
+py_binary(
+    name = "py_test_bin",
+    srcs = ["@pypi_pytest//:rules_python_wheel_entry_point_pytest"],
+    # rules_python permits this label even though it's not a source file in this repository.
+    main = "rules_python_wheel_entry_point_pytest.py",
+    deps = ["@pypi_pytest//:pkg"],
+)
+
+genrule(
+    name = "run_py_test_help",
+    outs = ["help.txt"],
+    cmd = "$(execpath py_test_bin) --help >$@",
+    tools = ["py_test_bin"],
+)
+
+assert_contains(
+    name = "test__run_py_test_help",
+    actual = "help.txt",
+    expected = "junit_logging (string):",
+)
+
+#################
+# Case 2:
+# The main property can be inferred from the target name.
+py_binary(
+    name = "main_from_name",
+    srcs = ["main_from_name.py"],
+)
+
+run_binary(
+    name = "run_main_from_name",
+    outs = ["main_from_name.out"],
+    args = ["$(execpath main_from_name.out)"],
+    tool = "main_from_name",
+)
+
+assert_contains(
+    name = "test__main_from_name",
+    actual = "main_from_name.out",
+    expected = "running main_from_name",
+)

py/tests/py-binary/main_from_name.py

@@ -0,0 +1,5 @@
+import sys
+
+if __name__ == '__main__':
+    with open(sys.argv[1], 'w') as f:
+        f.write("running main_from_name")