From 956e1bfcf342090bef931761782c270d2e8614e9 Mon Sep 17 00:00:00 2001
From: Alex Eagle <alex.eagle@robinhood.com>
Date: Wed, 8 Sep 2021 11:01:05 -0700
Subject: [PATCH] Allow for requirements files to differ per platform

As a common example, we need a compiled requirements file for linux that differs from mac os
---
 python/pip.bzl                        |  2 +-
 python/pip_install/pip_repository.bzl | 14 +++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/python/pip.bzl b/python/pip.bzl
index 785156f4b7..a00353133e 100644
--- a/python/pip.bzl
+++ b/python/pip.bzl
@@ -16,7 +16,7 @@
 load("//python/pip_install:pip_repository.bzl", "pip_repository")
 load("//python/pip_install:repositories.bzl", "pip_install_dependencies")
 
-def pip_install(requirements, name = "pip", **kwargs):
+def pip_install(requirements = None, name = "pip", **kwargs):
     """Imports a `requirements.txt` file and generates a new `requirements.bzl` file.
 
     This is used via the `WORKSPACE` pattern:
diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
index 71d5f662e9..c854c58947 100644
--- a/python/pip_install/pip_repository.bzl
+++ b/python/pip_install/pip_repository.bzl
@@ -115,12 +115,21 @@ def _pip_repository_impl(rctx):
             str(rctx.attr.timeout),
         ]
     else:
+        os = rctx.os.name.lower()
+        requirements_txt = rctx.attr.requirements
+
+        if os in rctx.attr.platform_requirements:
+            requirements_txt = Label(rctx.attr.platform_requirements[os])
+
+        if requirements_txt == None:
+            fail("You must set the requirements attribute, or platform_requirements attribute must include key " + os)
+
         args = [
             python_interpreter,
             "-m",
             "python.pip_install.extract_wheels",
             "--requirements",
-            rctx.path(rctx.attr.requirements),
+            rctx.path(requirements_txt),
         ]
 
     args += ["--repo", rctx.attr.name]
@@ -204,6 +213,9 @@ pip_repository_attrs = {
         default = False,
         doc = "Create the repository in incremental mode.",
     ),
+    "platform_requirements": attr.string_dict(
+        doc = "Override the requirements attribute when running on a matching host platform",
+    ),
     "requirements": attr.label(
         allow_single_file = True,
         doc = "A 'requirements.txt' pip requirements file.",