feat: Rule for C extensions

docs/pyproject.toml

@@ -13,5 +13,6 @@ dependencies = [
     "absl-py",
     "typing-extensions",
     "sphinx-reredirects",
-    "pefile"
+    "pefile",
+    "pyelftools",
 ]

docs/requirements.txt

@@ -236,6 +236,10 @@ pefile==2024.8.26 \
     --hash=sha256:3ff6c5d8b43e8c37bb6e6dd5085658d658a7a0bdcd20b6a07b1fcfc1c4e9d632 \
     --hash=sha256:76f8b485dcd3b1bb8166f1128d395fa3d87af26360c2358fb75b80019b957c6f
     # via rules-python-docs (docs/pyproject.toml)
+pyelftools==0.32 \
+    --hash=sha256:013df952a006db5e138b1edf6d8a68ecc50630adbd0d83a2d41e7f846163d738 \
+    --hash=sha256:6de90ee7b8263e740c8715a925382d4099b354f29ac48ea40d840cf7aa14ace5
+    # via rules-python-docs (docs/pyproject.toml)
 pygments==2.19.2 \
     --hash=sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887 \
     --hash=sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b

python/cc/py_extension.bzl

@@ -0,0 +1,8 @@
+"""Public API for py_extension."""
+
+load(
+    "//python/private/cc:py_extension_macro.bzl",
+    _py_extension = "py_extension",
+)
+
+py_extension = _py_extension

python/private/cc/py_extension_macro.bzl

@@ -0,0 +1,39 @@
+"""Wrapper macro for the py_extension rule."""
+
+load("@rules_cc//cc:cc_shared_library.bzl", "cc_shared_library")
+load("//python/private:util.bzl", "add_tag")
+load(
+    ":py_extension_rule.bzl",
+    _py_extension = "py_extension",
+    ##_py_extension_csl_rule = "py_extension_csl",
+)
+
+def py_extension(**kwargs):
+    """A macro that calls the py_extension rule and adds a tag.
+
+    Args:
+        **kwargs: Additional arguments to pass to the rule.
+    """
+    add_tag(kwargs, "@rules_python//python/cc:py_extension")
+
+    use_csl = kwargs.pop("use_csl", False)
+    if use_csl:
+        _py_extension_csl(**kwargs)
+    else:
+        if "libc" not in kwargs:
+            kwargs["libc"] = select({
+                "//conditions:default": "glibc",
+                "@rules_python//python/config_settings:_is_py_linux_libc_glibc": "glibc",
+                "@rules_python//python/config_settings:_is_py_linux_libc_musl": "musl",
+            })
+        _py_extension(**kwargs)
+
+def _py_extension_csl(*, name, module_name = None, **kwargs):
+    if not module_name:
+        module_name = name
+
+    cc_shared_library(
+        name = name,
+        shared_lib_name = module_name + ".so",
+        **kwargs
+    )

python/private/cc/py_extension_rule.bzl

@@ -0,0 +1,297 @@
+"""Implementation of the py_extension rule."""
+
+load("@rules_cc//cc/common:cc_common.bzl", "cc_common")
+load("@rules_cc//cc/common:cc_info.bzl", "CcInfo")
+load("//python:versions.bzl", "PLATFORMS")
+load("//python/private:attr_builders.bzl", "attrb")
+load("//python/private:attributes.bzl", "COMMON_ATTRS")
+load("//python/private:builders.bzl", "builders")
+load("//python/private:py_info.bzl", "PyInfo")
+load("//python/private:py_internal.bzl", "py_internal")
+load("//python/private:reexports.bzl", "BuiltinPyInfo")
+load("//python/private:rule_builders.bzl", "ruleb")
+load("//python/private:toolchain_types.bzl", "PY_CC_TOOLCHAIN_TYPE")
+
+def _py_extension_impl(ctx):
+    module_name = ctx.attr.module_name or ctx.label.name
+    repo_name = ctx.label.workspace_name or ctx.workspace_name
+    import_path = repo_name
+    if ctx.label.package:
+        import_path = repo_name + "/" + ctx.label.package
+    cc_toolchain = ctx.toolchains["@bazel_tools//tools/cpp:toolchain_type"].cc
+    feature_configuration = cc_common.configure_features(
+        ctx = ctx,
+        cc_toolchain = cc_toolchain,
+    )
+
+    # Collect CcInfo from all deps for compilation
+    static_deps_infos = [dep[CcInfo] for dep in ctx.attr.static_deps]
+    dynamic_deps_infos = [dep[CcSharedLibraryInfo] for dep in ctx.attr.dynamic_deps]
+    external_deps_infos = [dep[CcInfo] for dep in ctx.attr.external_deps]
+
+    # Static deps are linked directly into the .so
+    static_cc_info = cc_common.merge_cc_infos(
+        cc_infos = static_deps_infos,
+    )
+
+    # Dynamic deps are linked as shared libraries
+    linker_inputs = [dep.linker_input for dep in dynamic_deps_infos]
+    dynamic_linking_context = cc_common.create_linking_context(
+        linker_inputs = depset(linker_inputs),
+    )
+
+    user_link_flags = []
+    user_link_flags.append("-Wl,--export-dynamic-symbol=PyInit_{module_name}".format(
+        module_name = module_name,
+    ))
+
+    # The PyInit symbol looks unused, so the linker optimizes it away. Telling it
+    # to treat it as undefined causes it to be retained.
+    user_link_flags.append("-Wl,--undefined=PyInit_{module_name}".format(
+        module_name = module_name,
+    ))
+
+    if ctx.attr.external_deps:
+        user_link_flags.append("-Wl,--allow-shlib-undefined")
+
+    ext = _get_extension(cc_toolchain)
+    use_py_limited_api = bool(ctx.attr.py_limited_api)
+    if use_py_limited_api:
+        output_filename = "{module_name}.abi3.{ext}".format(
+            module_name = module_name,
+            ext = ext,
+        )
+    else:
+        py_toolchain = ctx.toolchains[PY_CC_TOOLCHAIN_TYPE]
+        py_cc_toolchain = py_toolchain.py_cc_toolchain
+        platform_tag = _get_platform(ctx)
+        output_filename = "{module_name}.{abi_tag}-{platform}.{ext}".format(
+            module_name = module_name,
+            abi_tag = py_cc_toolchain.abi_tag,
+            platform = platform_tag,
+            ext = ext,
+        )
+    py_dso = ctx.actions.declare_file(output_filename)
+
+    static_linking_context = static_cc_info.linking_context
+    linking_contexts = [
+        static_linking_context,
+        dynamic_linking_context,
+    ]
+
+    # Add target-level linkopts last so users can override.
+    user_link_flags.extend(ctx.attr.linkopts)
+
+    # todo: add linker script to hide symbols by default
+    # py_internal allows using some private apis, which may or may not be needed.
+    # based upon cc_shared_library.bzl
+    cc_linking_outputs = py_internal.link(
+        actions = ctx.actions,
+        feature_configuration = feature_configuration,
+        cc_toolchain = cc_toolchain,
+        linking_contexts = linking_contexts,
+        user_link_flags = user_link_flags,
+        # todo: add additional_inputs
+        name = ctx.label.name,
+        output_type = "dynamic_library",
+        main_output = py_dso,
+        # todo: maybe variables_extension
+        # todo: maybe additional_outputs
+    )
+
+    # Propagate CcInfo from dynamic and external deps, but not static ones.
+    dynamic_cc_info = CcInfo(linking_context = dynamic_linking_context)
+    propagated_cc_info = cc_common.merge_cc_infos(
+        cc_infos = [dynamic_cc_info] + external_deps_infos,
+    )
+
+    runfiles_builder = builders.RunfilesBuilder()
+    runfiles_builder.add(py_dso)
+    runfiles_builder.add_targets(ctx.attr.static_deps)
+    runfiles_builder.add_targets(ctx.attr.dynamic_deps)
+    runfiles_builder.add_targets(ctx.attr.external_deps)
+    runfiles = runfiles_builder.build(ctx)
+
+    return [
+        DefaultInfo(
+            files = depset([py_dso]),
+            runfiles = runfiles,
+        ),
+        PyInfo(
+            transitive_sources = depset([py_dso]),
+            imports = depset([import_path]),
+        ),
+        propagated_cc_info,
+    ]
+
+
+PY_EXTENSION_ATTRS = COMMON_ATTRS | {
+    "dynamic_deps": lambda: attrb.LabelList(
+        providers = [CcSharedLibraryInfo],
+        doc = "cc_shared_library targets to be dynamically linked.",
+        default = [],
+    ),
+    "external_deps": lambda: attrb.LabelList(
+        providers = [CcInfo],
+        doc = "cc_library targets with external linkage.",
+        default = [],
+    ),
+    "static_deps": lambda: attrb.LabelList(
+        providers = [CcInfo],
+        doc = "cc_library targets to be statically and privately linked.",
+        default = [],
+    ),
+    "copts": lambda: attrb.StringList(),
+    "libc": lambda: attrb.String(default = "glibc"),
+    "linkopts": lambda: attrb.StringList(),
+    "module_name": lambda: attrb.String(),
+    "py_limited_api": lambda: attrb.String(
+        doc = """
+The minimum Python version to target for the Limited API (e.g., '3.8').
+
+If set to a version string (e.g., '3.8') instead of '' (empty string):
+  - Configures the output filename to use the simple '.abi3' suffix
+    (e.g., 'ext.abi3.so').
+
+Note: Since the py_extension rule only links pre-compiled libraries,
+you must manually add the preprocessor macro to the cc_library targets
+that compile your C/C++ sources, for example:
+    cc_library(
+        name = "my_impl",
+        srcs = ["my_code.c"],
+        defines = ["Py_LIMITED_API=0x03080000"],
+        ...
+    )
+
+Set to '' (the default) or None to build a standard, version-specific
+extension.
+""",
+        default = "",
+    ),
+    "_constraints": lambda: attrb.LabelList(
+        default = sorted({
+            c: None
+            for info in PLATFORMS.values()
+            for c in info.compatible_with
+        }.keys()),
+    ),
+}
+
+def create_py_extension_rule_builder(**kwargs):
+    """Create a rule builder for a py_extension."""
+    builder = ruleb.Rule(
+        implementation = _py_extension_impl,
+        attrs = PY_EXTENSION_ATTRS,
+        provides = [PyInfo, CcInfo],
+        toolchains = [
+            ruleb.ToolchainType(PY_CC_TOOLCHAIN_TYPE),
+            ruleb.ToolchainType("@bazel_tools//tools/cpp:toolchain_type"),
+        ],
+        fragments = ["cpp"],
+        **kwargs
+    )
+    return builder
+
+py_extension = create_py_extension_rule_builder().build()
+
+def _get_extension(cc_toolchain):
+    """
+    Derives the appropriate file extension from the C++ toolchain.
+
+    Args:
+        cc_toolchain: The CcToolchainInfo provider (usually obtained via
+          ctx.toolchains["@bazel_tools//tools/cpp:toolchain_type"].cc)
+
+    Returns:
+        The extension, e.g. "so" or "pyd"
+    """
+
+    # Windows uses .pyd; Unix (Linux/macOS) uses .so for Python modules
+    target_name = cc_toolchain.target_gnu_system_name
+    is_windows = "windows" in target_name or "mingw" in target_name or "msvc" in target_name
+    ext = "pyd" if is_windows else "so"
+    return ext
+
+def _derive_pep3149_tag(platform, info):
+    # platform is the triplet, e.g. "x86_64-unknown-linux-gnu"
+    p, _, _ = platform.partition("-freethreaded")
+    parts = p.split("-")
+    triplet_arch = parts[0]
+
+    if info.os_name == "windows":
+        if triplet_arch == "x86_64":
+            return "win_amd64"
+        elif triplet_arch == "aarch64":
+            return "win_arm64"
+        else:
+            return "win32"
+    elif info.os_name == "osx":
+        return "darwin"
+    elif info.os_name == "linux":
+        abi = "musl" if p.endswith("-musl") else "gnu"
+        return "{}-linux-{}".format(triplet_arch, abi)
+    else:
+        return triplet_arch
+
+def _get_platform_from_constraints(ctx):
+    # Build a map of Label to ConstraintValueInfo from _constraints
+    constraints_map = {}
+    for c in ctx.attr._constraints:
+        if platform_common.ConstraintValueInfo in c:
+            constraints_map[c.label] = c[platform_common.ConstraintValueInfo]
+
+    # Resolve the target's libc to its config_setting label string
+    target_libc_setting = None
+    if ctx.attr.libc == "musl":
+        target_libc_setting = str(Label("//python/config_settings:_is_py_linux_libc_musl"))
+    elif ctx.attr.libc == "glibc":
+        target_libc_setting = str(Label("//python/config_settings:_is_py_linux_libc_glibc"))
+
+    # Find the matching platform in PLATFORMS
+    for platform, info in PLATFORMS.items():
+        # Check if all compatible_with constraints are satisfied
+        match = True
+        for c_str in info.compatible_with:
+            c_label = Label(c_str)
+            if c_label in constraints_map:
+                c_val = constraints_map[c_label]
+                if not ctx.target_platform_has_constraint(c_val):
+                    match = False
+                    break
+            else:
+                match = False
+                break
+
+        if match:
+            # Additional check for Linux libc consistency using target_settings
+            if info.os_name == "linux" and target_libc_setting:
+                if target_libc_setting not in info.target_settings:
+                    match = False
+
+        if match:
+            return _derive_pep3149_tag(platform, info)
+
+    return None
+
+def _get_platform(ctx):
+    """Derives the PEP 3149 platform tag from the target constraints.
+
+    Args:
+        ctx: The rule context.
+
+    Returns:
+        The platform tag, e.g. "x86_64-linux-gnu" or "win_amd64"
+    """
+    platform_tag = _get_platform_from_constraints(ctx)
+    if platform_tag:
+        return platform_tag
+
+    fail(
+        """
+ERROR: Unsupported target platform for {self}.
+  The target platform's constraints do not match any supported platform
+  in rules_python's central registry (python/versions.bzl).
+  Please ensure your target platform is configured correctly.""".format(
+            self = ctx.label,
+        ),
+    )

python/private/py_cc_toolchain_info.bzl

@@ -17,6 +17,11 @@
 PyCcToolchainInfo = provider(
     doc = "C/C++ information about the Python runtime.",
     fields = {
+        "abi_tag": """\
+:type: str
+
+The ABI tag for extension modules, e.g. 'cpython-311' or 'cpython-313t'.
+""",
         "headers": """\
 :type: struct