From 5dd61aa51d148e3583576338d79292c2c208be5b Mon Sep 17 00:00:00 2001
From: Nicolas capens <ncapens@nvidia.com>
Date: Wed, 7 Jun 2023 15:05:39 -0400
Subject: [PATCH] Support specifying the wheel platform

Calling `python -m build --wheel -C--build-option=-Pwindows` now creates
dist/warp_lang-0.9.0-py3-none-win_amd64.whl, even if warp/bin/ contains
libraries other than .dll.

Likewise -C--build-option=-Plinux and -C--build-option=-Pmacos package
just the .so and .dylib libraries, respectively. This enables creating
three platform-specific wheels from a single artifact archive containing
binary libraries for all three platforms.

Executing just `python -m build --wheel` is still supported and will
select a single platform, based on the libraries it detects. This is
useful for creating a wheel from an artifact archive containing
libraries for a single platform.
---
 .gitignore   |   2 +
 PACKAGING.md |  11 +++--
 setup.py     | 123 +++++++++++++++++++++++++++++++++++++++++++--------
 3 files changed, 114 insertions(+), 22 deletions(-)

diff --git a/.gitignore b/.gitignore
index dd5cc80a6..90e215d11 100644
--- a/.gitignore
+++ b/.gitignore
@@ -26,3 +26,5 @@ build/lib/
 .venv/
 exts/omni.warp/config/extension.gen.toml
 /external/llvm-project
+/build
+/dist
diff --git a/PACKAGING.md b/PACKAGING.md
index 5747f32ea..10ef78bd2 100644
--- a/PACKAGING.md
+++ b/PACKAGING.md
@@ -71,9 +71,14 @@ Release Steps:
 
 2) Run `cd warp`
 
-2) Run `python -m build`
-
-3) Run `python -m twine upload dist/*`
+3) Run
+    ```bash
+    python -m build --wheel -C--build-option=-Pwindows &&
+    python -m build --wheel -C--build-option=-Plinux &&
+    python -m build --wheel -C--build-option=-Pmacos
+    ```
+
+4) Run `python -m twine upload dist/*`
 
     * user: `__token__`
     * pass: `(your token string from pypi)`
diff --git a/setup.py b/setup.py
index c4c318c5c..097de7edf 100644
--- a/setup.py
+++ b/setup.py
@@ -1,18 +1,99 @@
 import setuptools
 import os
+import shutil
+import argparse
+
+from typing import NamedTuple
 from wheel.bdist_wheel import bdist_wheel
 
+# Parse --build-option arguments meant for the bdist_wheel command. We have to parse these
+# ourselves because when bdist_wheel runs it's too late to select a subset of libraries for package_data.
+parser = argparse.ArgumentParser(exit_on_error=False)
+parser.add_argument("command")
+parser.add_argument("--platform", "-P", type=str, default="", help="Wheel platform: windows|linux|macos")
+args = parser.parse_known_args()[0]
+
+
+class Platform(NamedTuple):
+    name: str
+    fancy_name: str
+    extension: str
+    tag: str
+
 
-def get_warp_platform():
+platforms = [
+    Platform("windows", "Windows", ".dll", "win_amd64"),
+    Platform("linux", "Linux", ".so", "manylinux2014_x86_64"),
+    Platform("macos", "macOS", ".dylib", "macosx_10_13_x86_64"),
+]
+
+
+def detect_warp_platforms():
+    detected_platforms = set()
     for filename in os.listdir("warp/bin"):
-        if os.path.splitext(filename)[1] == ".dll":
-            return "win_amd64"
-        if os.path.splitext(filename)[1] == ".so":
-            return "manylinux2014_x86_64"
-        if os.path.splitext(filename)[1] == ".dylib":
-            return "macosx_10_13_x86_64"
+        for p in platforms:
+            if os.path.splitext(filename)[1] == p.extension:
+                detected_platforms.add(p)
+
+    if len(detected_platforms) == 0:
+        raise Exception("No libraries found in warp/bin. Please run build_lib.py first.")
+
+    return detected_platforms
+
+
+wheel_platform = None
+
+if args.command == "bdist_wheel":
+    detected_platforms = detect_warp_platforms()
+
+    if args.platform != "":
+        for p in platforms:
+            if args.platform == p.name or args.platform == p.fancy_name:
+                wheel_platform = p
+                print(f"Platform argument specified for building {p.fancy_name} wheel")
+                break
+
+        if wheel_platform is None:
+            print(f"Platform argument '{args.platform}' not recognized")
+        elif wheel_platform not in detected_platforms:
+            print(f"No libraries found for {wheel_platform.fancy_name}")
+            print(f"Falling back to auto-detection")
+            wheel_platform = None
+
+    if wheel_platform is None:
+        if len(detected_platforms) > 1:
+            print("Libraries for multiple platforms were detected. Picking the first one.")
+            print("Run `python -m build --wheel -C--build-option=-P[windows|linux|macos]` to select a specific one.")
+        wheel_platform = next(iter(detected_platforms))
 
-    raise Exception("No libraries found in warp/bin")
+    print("Creating Warp wheel for " + wheel_platform.fancy_name)
+
+
+# Binary wheel distribution builds assume that the platform you're building on will be the platform
+# of the package. This class overrides the platform tag.
+# https://packaging.python.org/en/latest/specifications/platform-compatibility-tags
+class WarpBDistWheel(bdist_wheel):
+    # Even though we parse the platform argument ourselves, we need to declare it here as well so
+    # setuptools.Command can validate the command line options.
+    user_options = bdist_wheel.user_options + [
+        ("platform=", "P", "Wheel platform: windows|linux|macos"),
+    ]
+
+    def initialize_options(self):
+        super().initialize_options()
+        self.platform = ""
+
+    def get_tag(self):
+        # The wheel's complete tag format is {python tag}-{abi tag}-{platform tag}.
+        return "py3", "none", wheel_platform.tag
+
+    def run(self):
+        super().run()
+
+        # Clean up so we can re-invoke `py -m build --wheel -C--build-option=--platform=...`
+        # See https://github.com/pypa/setuptools/issues/1871 for details.
+        shutil.rmtree("./build", ignore_errors=True)
+        shutil.rmtree("./warp_lang.egg-info", ignore_errors=True)
 
 
 # Distributions are identified as non-pure (i.e. containing non-Python code, or binaries) if the
@@ -23,17 +104,19 @@ def has_ext_modules(self):
         return True
 
 
-# Binary wheel distribution builds assume that the platform you're building on will be the platform
-# of the package. This factory function provides a class which overrides the platform tag.
-# https://packaging.python.org/en/latest/specifications/platform-compatibility-tags
-def bdist_factory(platform_tag: str):
-    class warp_bdist(bdist_wheel):
-        def get_tag(self, *args, **kws):
-            # The tag format is {python tag}-{abi tag}-{platform tag}.
-            return "py3", "none", platform_tag
+def get_warp_libraries(extension):
+    libraries = []
+    for filename in os.listdir("warp/bin"):
+        if os.path.splitext(filename)[1] == extension:
+            libraries.append("bin/" + filename)
+
+    return libraries
 
-    return warp_bdist
 
+if wheel_platform is not None:
+    warp_binary_libraries = get_warp_libraries(wheel_platform.extension)
+else:
+    warp_binary_libraries = []  # Not needed during egg_info command
 
 setuptools.setup(
     name="warp-lang",
@@ -57,8 +140,8 @@ def get_tag(self, *args, **kws):
             "native/clang/*.cpp",
             "native/nanovdb/*.h",
             "tests/assets/*",
-            "bin/*",
         ]
+        + warp_binary_libraries,
     },
     classifiers=[
         "Programming Language :: Python :: 3.7",
@@ -70,7 +153,9 @@ def get_tag(self, *args, **kws):
         "Operating System :: OS Independent",
     ],
     distclass=BinaryDistribution,
-    cmdclass={"bdist_wheel": bdist_factory(get_warp_platform())},
+    cmdclass={
+        "bdist_wheel": WarpBDistWheel,
+    },
     install_requires=["numpy"],
     python_requires=">=3.7",
 )