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,107 @@
+"""Macro for creating Python extensions."""
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_shared_library.bzl", "cc_shared_library")
+load("//python/private:util.bzl", "add_tag", "copy_propagating_kwargs")
+load(":py_extension_rule.bzl", "py_extension_wrapper")
+
+def py_extension(
+        name,
+        srcs = None,
+        hdrs = None,
+        copts = None,
+        defines = None,
+        deps = None,
+        dynamic_deps = None,
+        exports_filter = None,
+        user_link_flags = None,
+        visibility = None,
+        data = None,
+        **kwargs):
+    """Creates a Python extension module.
+
+    By default, extensions are created within their workspace package directory
+    (e.g., `pkg/ext.so`) and imported using standard Python package paths
+    (e.g., `from pkg import ext`).
+
+    To customize import path behavior:
+    - `imports`: Pass `imports = ["..."]` to append custom search directories to
+      `sys.path` (matching `py_library`).
+    - `module_name`: Pass `module_name = "custom_name"` to override the base module
+      filename.
+
+    Args:
+        name: Target name.
+        srcs: Optional C/C++ source files to compile directly for this extension.
+        hdrs: Optional header files for the srcs.
+        copts: Optional compiler flags for srcs.
+        defines: Optional preprocessor defines for srcs.
+        deps: cc_library targets to statically link into the extension.
+        dynamic_deps: cc_shared_library targets to dynamically link.
+        exports_filter: Filter for exported symbols passed to cc_shared_library.
+        user_link_flags: Additional link flags passed to cc_shared_library.
+        visibility: Target visibility.
+        data: Optional list of files or targets needed by this extension at runtime.
+        **kwargs: Additional arguments passed to the underlying wrapper rule.
+    """
+    add_tag(kwargs, "@rules_python//python/cc:py_extension")
+
+    csl_deps = []
+
+    # 1. Handle user-supplied static deps
+    if deps:
+        csl_deps.extend(deps)
+
+    # 2. If srcs or hdrs are specified, create an implicit cc_library for them
+    if srcs or hdrs:
+        impl_lib_name = "_" + name + "_impl"
+        cc_library(
+            name = impl_lib_name,
+            srcs = srcs,
+            hdrs = hdrs,
+            copts = (copts or []) + ["-fPIC"],
+            defines = defines,
+            deps = ["@rules_python//python/cc:current_py_cc_headers"],
+            visibility = ["//visibility:private"],
+            **copy_propagating_kwargs(kwargs)
+        )
+        csl_deps.append(":" + impl_lib_name)
+
+    # 3. If no static deps or sources were specified, use empty target for CSL requirement
+    if not csl_deps:
+        csl_deps.append("//python/private/cc:empty")
+
+    # 4. Create the underlying cc_shared_library
+    csl_name = "_" + name + "_csl"
+    csl_kwargs = copy_propagating_kwargs(kwargs)
+    if exports_filter:
+        csl_kwargs["exports_filter"] = exports_filter
+    if user_link_flags:
+        csl_kwargs["user_link_flags"] = user_link_flags
+
+    cc_shared_library(
+        name = csl_name,
+        deps = csl_deps,
+        dynamic_deps = dynamic_deps,
+        visibility = ["//visibility:private"],
+        **csl_kwargs
+    )
+
+    # 5. Select default libc constraint if not provided
+    if "libc" not in kwargs:
+        kwargs["libc"] = select({
+            "@rules_python//python/config_settings:_is_py_linux_libc_glibc": "glibc",
+            "@rules_python//python/config_settings:_is_py_linux_libc_musl": "musl",
+            "//conditions:default": "glibc",
+        })
+
+    if data != None:
+        kwargs["data"] = data
+
+    # 6. Wrap with py_extension_wrapper for PEP 3149 naming & PyInfo
+    py_extension_wrapper(
+        name = name,
+        src = ":" + csl_name,
+        visibility = visibility,
+        **kwargs
+    )

python/private/cc/py_extension_rule.bzl

@@ -0,0 +1,149 @@
+"""Implementation of the _py_extension_wrapper rule."""
+
+load("@bazel_skylib//lib:dicts.bzl", "dicts")
+load("@rules_cc//cc/common:cc_shared_library_info.bzl", "CcSharedLibraryInfo")
+load("//python/private:attr_builders.bzl", "attrb")
+load("//python/private:attributes.bzl", "COMMON_ATTRS", "IMPORTS_ATTRS")
+load("//python/private:builders.bzl", "builders")
+load("//python/private:py_info.bzl", "PyInfo")
+load("//python/private:rule_builders.bzl", "ruleb")
+load("//python/private:toolchain_types.bzl", "PY_CC_TOOLCHAIN_TYPE")
+
+def _py_extension_wrapper_impl(ctx):
+    module_name = ctx.attr.module_name or ctx.label.name
+
+    cc_toolchain = ctx.toolchains["@bazel_tools//tools/cpp:toolchain_type"].cc
+    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)
+
+    # Symlink the cc_shared_library output to the PEP 3149 / abi3 filename
+    csl_target = ctx.attr.src
+    csl_file = csl_target[DefaultInfo].files.to_list()[0]
+    ctx.actions.symlink(
+        output = py_dso,
+        target_file = csl_file,
+    )
+
+    runfiles_builder = builders.RunfilesBuilder()
+    runfiles_builder.add(py_dso)
+    runfiles_builder.add(ctx.files.data)
+    runfiles_builder.add_targets(ctx.attr.data)
+    runfiles_builder.add(csl_target[DefaultInfo].default_runfiles)
+    runfiles = runfiles_builder.build(ctx)
+
+    # Resolve imports paths relative to the target package and repository:
+    # 1. Default (imports = []): No extra search paths are added to sys.path,
+    #    enforcing clean package-qualified imports (e.g. `from foo.bar import ext`).
+    # 2. Relative paths (e.g. imports = ["."]): Resolved relative to `repo_name/package_name`
+    #    so passing `imports = ["."]` adds the target's package directory to sys.path.
+    # 3. Absolute paths (starting with "/"): Stripped of leading "/" and resolved relative to runfiles root.
+    imports_list = []
+    repo_name = ctx.label.workspace_name or ctx.workspace_name
+    for path in ctx.attr.imports:
+        if path.startswith("/"):
+            imports_list.append(path[1:])
+        else:
+            pkg = ctx.label.package
+            full_path = "{}/{}".format(pkg, path) if pkg else path
+            if repo_name:
+                full_path = "{}/{}".format(repo_name, full_path)
+            imports_list.append(full_path)
+
+    return [
+        DefaultInfo(
+            files = depset([py_dso]),
+            runfiles = runfiles,
+        ),
+        PyInfo(
+            transitive_sources = depset([py_dso]),
+            imports = depset(imports_list),
+        ),
+    ]
+
+PY_EXTENSION_WRAPPER_ATTRS = dicts.add(
+    COMMON_ATTRS,
+    IMPORTS_ATTRS,
+    {
+        "libc": lambda: attrb.String(default = "glibc"),
+        "module_name": lambda: attrb.String(),
+        "py_limited_api": lambda: attrb.String(
+            default = "",
+        ),
+        "src": lambda: attrb.Label(
+            mandatory = True,
+            providers = [CcSharedLibraryInfo],
+            doc = "The cc_shared_library target to wrap.",
+        ),
+    },
+)
+
+def create_py_extension_wrapper_rule_builder(**kwargs):
+    """Create a rule builder for the wrapper."""
+    builder = ruleb.Rule(
+        implementation = _py_extension_wrapper_impl,
+        attrs = PY_EXTENSION_WRAPPER_ATTRS,
+        provides = [PyInfo],
+        toolchains = [
+            ruleb.ToolchainType(PY_CC_TOOLCHAIN_TYPE),
+            ruleb.ToolchainType("@bazel_tools//tools/cpp:toolchain_type"),
+        ],
+        fragments = ["cpp"],
+        **kwargs
+    )
+    return builder
+
+py_extension_wrapper = create_py_extension_wrapper_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 _get_platform(ctx):
+    """Derives the PEP 3149 platform tag from the active Python C++ toolchain.
+
+    Args:
+        ctx: The rule context.
+
+    Returns:
+        The platform tag, e.g. "x86_64-linux-gnu" or "win_amd64"
+    """
+    py_toolchain = ctx.toolchains[PY_CC_TOOLCHAIN_TYPE]
+    py_cc_toolchain = py_toolchain.py_cc_toolchain
+    if hasattr(py_cc_toolchain, "platform_tag") and py_cc_toolchain.platform_tag:
+        return py_cc_toolchain.platform_tag
+
+    fail(
+        "ERROR: Unable to resolve platform_tag from Python C++ toolchain for {self}. " +
+        "Please ensure the active py_cc_toolchain provides a non-empty platform_tag.",
+    )

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
 
@@ -91,6 +96,11 @@ If available, information about C libraries, struct with fields:
     considered private and should be forward along as-is (this better allows
     e.g. `:current_py_cc_headers` to act as the underlying headers target it
     represents).
+""",
+        "platform_tag": """\
+:type: str | None
+
+The PEP 3149 / PEP 425 platform tag for extension modules, e.g. 'x86_64-linux-gnu', 'darwin', or 'win_amd64'.
 """,
         "python_version": """
 :type: str

python/private/py_cc_toolchain_macro.bzl

@@ -30,4 +30,23 @@ def py_cc_toolchain(**kwargs):
 
     #  This tag is added to easily identify usages through other macros.
     add_tag(kwargs, "@rules_python//python:py_cc_toolchain")
+
+    if "os" not in kwargs:
+        kwargs["os"] = select({
+            "@platforms//os:macos": "macos",
+            "@platforms//os:windows": "windows",
+            "//conditions:default": "linux",
+        })
+    if "cpu" not in kwargs:
+        kwargs["cpu"] = select({
+            "@platforms//cpu:aarch64": "aarch64",
+            "@platforms//cpu:x86_32": "x86_32",
+            "//conditions:default": "x86_64",
+        })
+    if "libc" not in kwargs:
+        kwargs["libc"] = select({
+            Label("//python/config_settings:_is_py_linux_libc_musl"): "musl",
+            "//conditions:default": "gnu",
+        })
+
     _py_cc_toolchain(**kwargs)