Expose C/C++ headers packaged in wheels

Summary:
Some `.whl`s include C/C++ headers that are used by dependents
to compile sources (e.g. numpy, torch).

This diff adds a couple ways to expose these headers for C/C++ rules
to consume, via a C/C++ preprocessor provider:
- `cxx_header_dirs`:  The use can specify the expected include dirs in
   the `.whl`
- `infer_cxx_header_dirs`: We'll find `include` dirs and add automatically
   export them

Reviewed By: christycylee

Differential Revision: D59762091

fbshipit-source-id: cc300e726172ea28d58bf282cbaf88dfe54441a7

decls/python_rules.bzl

@@ -191,6 +191,8 @@ prebuilt_python_library = prelude_rule(
         {
             "compile": attrs.bool(default = False),
             "contacts": attrs.list(attrs.string(), default = []),
+            "cxx_header_dirs": attrs.option(attrs.list(attrs.string()), default = None),
+            "infer_cxx_header_dirs": attrs.bool(default = False),
             "default_host_platform": attrs.option(attrs.configuration_label(), default = None),
             "ignore_compile_errors": attrs.bool(default = False),
             "licenses": attrs.list(attrs.source(), default = []),

python/prebuilt_python_library.bzl

@@ -15,6 +15,14 @@ load(
     "get_excluded",
     "get_roots",
 )
+load(
+    "@prelude//cxx:preprocessor.bzl",
+    "CPreprocessor",
+    "CPreprocessorArgs",
+    "cxx_inherited_preprocessor_infos",
+    "cxx_merge_cpreprocessors",
+    "format_system_include_arg",
+)
 load(
     "@prelude//linking:linkable_graph.bzl",
     "create_linkable_graph",
@@ -57,6 +65,13 @@ def prebuilt_python_library_impl(ctx: AnalysisContext) -> list[Provider]:
     )
     if ctx.attrs.strip_soabi_tags:
         cmd.add("--strip-soabi-tags")
+    inferred_cxx_header_dirs = None
+    if ctx.attrs.infer_cxx_header_dirs:
+        inferred_cxx_header_dirs = ctx.actions.declare_output("__cxx_header_dirs__.txt")
+        cmd.add(
+            "--cxx-header-dirs",
+            inferred_cxx_header_dirs.as_output(),
+        )
     ctx.actions.run(cmd, category = "py_extract_prebuilt_library")
     deps, shared_deps = gather_dep_libraries(ctx.attrs.deps)
     src_manifest = create_manifest_for_source_dir(ctx, "binary_src", extracted_src, exclude = "\\.pyc$")
@@ -142,4 +157,56 @@ def prebuilt_python_library_impl(ctx: AnalysisContext) -> list[Provider]:
         ),
     )
 
+    # If this prebuilt wheel contains headers, export them via a C++ provider.
+    pp_args = []
+    if ctx.attrs.cxx_header_dirs:
+        for header_dir in ctx.attrs.cxx_header_dirs:
+            pp_args.append(
+                format_system_include_arg(
+                    cmd_args(extracted_src.project(header_dir)),
+                    "clang",
+                ),
+            )
+    if inferred_cxx_header_dirs != None:
+        pp_argsfile = ctx.actions.declare_output("__cxx_header_dirs__.argsfile")
+
+        def write_argsfile(actions, header_dirs, output):
+            lines = []
+            for header_dir in header_dirs.read_string().splitlines():
+                lines.append(format_system_include_arg(
+                    cmd_args(extracted_src.project(header_dir)),
+                    "clang",
+                ))
+            actions.write(output, lines)
+
+        ctx.actions.dynamic_output(
+            dynamic = [inferred_cxx_header_dirs],
+            inputs = [],
+            outputs = [pp_argsfile.as_output()],
+            f = lambda ctx, artifacts, outputs: write_argsfile(
+                ctx.actions,
+                artifacts[inferred_cxx_header_dirs],
+                outputs[pp_argsfile],
+            ),
+        )
+        pp_args.append(
+            cmd_args(
+                pp_argsfile,
+                format = "@{}",
+                hidden = [extracted_src],
+            ),
+        )
+    if pp_args:
+        providers.append(cxx_merge_cpreprocessors(
+            ctx = ctx,
+            own = [
+                CPreprocessor(
+                    args = CPreprocessorArgs(
+                        args = pp_args,
+                    ),
+                ),
+            ],
+            xs = cxx_inherited_preprocessor_infos(ctx.attrs.deps),
+        ))
+
     return providers

python/tools/extract.py

@@ -104,6 +104,11 @@ def main() -> None:
     )
     parser.add_argument("--strip-soabi-tags", action="store_true")
     parser.add_argument("--entry-points", type=Path, help="The directory to write to")
+    parser.add_argument(
+        "--cxx-header-dirs",
+        type=Path,
+        help="A file to write out inferred C++ include dirs to",
+    )
     parser.add_argument(
         "--entry-points-manifest", type=Path, help="The directory to write to"
     )
@@ -118,6 +123,14 @@ def main() -> None:
         strip_soabi_tags=args.strip_soabi_tags,
     )
 
+    # Infer C++ header dirs.
+    if args.cxx_header_dirs is not None:
+        with open(args.cxx_header_dirs, mode="w") as f:
+            for root, dirs, _files in os.walk(args.output):
+                root = os.path.relpath(root, args.output)
+                if "include" in dirs:
+                    print(os.path.normpath(os.path.join(root, "include")), file=f)
+
     # Extract any "entry points" from the wheel, and generate scripts from them
     # (just like `pip install` would do).
     if args.entry_points is not None: