From 4f0db407cfb09a0e7cd5dbc5072f6e4efcf2df8a Mon Sep 17 00:00:00 2001
From: Ben Rogers <benrogers@mac.com>
Date: Thu, 3 Jul 2025 11:01:44 -0700
Subject: [PATCH] feat(fbcode_builder): Enable python `pex` archives to allow
 native library dependencies

Summary:
Fixes integration test failure due to lack of a functioning
python-psutil (on Linux, this requires the inclusion of native code
which is not permitted in a zipapp).

X-link: https://github.com/facebook/sapling/pull/1100

Reviewed By: zzl0

Differential Revision: D77674102

Pulled By: quark-zju

fbshipit-source-id: 11ac197d8c4082eaf16e0e28bc1a45c67f7dbb07
---
 .../fbcode_builder/CMake/FBPythonBinary.cmake |  6 ++++
 .../fbcode_builder/CMake/make_fbpy_archive.py | 32 +++++++++++++++++++
 build/fbcode_builder/getdeps.py               | 30 ++++++++++-------
 3 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/build/fbcode_builder/CMake/FBPythonBinary.cmake b/build/fbcode_builder/CMake/FBPythonBinary.cmake
index f91ebaf32..2be86fcb3 100644
--- a/build/fbcode_builder/CMake/FBPythonBinary.cmake
+++ b/build/fbcode_builder/CMake/FBPythonBinary.cmake
@@ -236,6 +236,7 @@ function(add_fb_python_unittest TARGET)
   set(
     one_value_args
     WORKING_DIRECTORY BASE_DIR NAMESPACE TEST_LIST DISCOVERY_TIMEOUT
+    WORKING_DIRECTORY BASE_DIR NAMESPACE TEST_LIST DISCOVERY_TIMEOUT TYPE
   )
   fb_cmake_parse_args(
     ARG "" "${one_value_args}" "${multi_value_args}" "${ARGN}"
@@ -286,8 +287,13 @@ function(add_fb_python_unittest TARGET)
   list(APPEND ARG_SOURCES "${FB_PY_TEST_MAIN}=__main__.py")
   list(APPEND ARG_SOURCES "${test_modules_path}=__test_modules__.py")
 
+  if(NOT DEFINED ARG_TYPE)
+    set(ARG_TYPE "zipapp")
+  endif()
+
   add_fb_python_executable(
     "${TARGET}"
+    TYPE "${ARG_TYPE}"
     NAMESPACE "${ARG_NAMESPACE}"
     BASE_DIR "${ARG_BASE_DIR}"
     SOURCES ${ARG_SOURCES}
diff --git a/build/fbcode_builder/CMake/make_fbpy_archive.py b/build/fbcode_builder/CMake/make_fbpy_archive.py
index 3724feb21..70d3c426e 100755
--- a/build/fbcode_builder/CMake/make_fbpy_archive.py
+++ b/build/fbcode_builder/CMake/make_fbpy_archive.py
@@ -7,6 +7,7 @@ import collections
 import errno
 import os
 import shutil
+import subprocess
 import sys
 import tempfile
 import zipapp
@@ -123,6 +124,36 @@ def populate_install_tree(inst_dir, path_map):
             pass
 
 
+def build_pex(args, path_map):
+    """Create a self executing python binary using the PEX tool
+
+    This type of Python binary is more complex as it requires a third-party tool,
+    but it does support native language extensions (.so/.dll files).
+    """
+    dest_dir = os.path.dirname(args.output)
+    with tempfile.TemporaryDirectory(prefix="make_fbpy.", dir=dest_dir) as tmpdir:
+        inst_dir = os.path.join(tmpdir, "tree")
+        populate_install_tree(inst_dir, path_map)
+
+        if os.path.exists(os.path.join(inst_dir, "__main__.py")):
+            os.rename(
+                os.path.join(inst_dir, "__main__.py"),
+                os.path.join(inst_dir, "main.py"),
+            )
+            args.main = "main"
+
+        tmp_output = os.path.abspath(os.path.join(tmpdir, "output.exe"))
+        subprocess.check_call(
+            ["pex"]
+            + ["--output-file", tmp_output]
+            + ["--python", args.python]
+            + ["--sources-directory", inst_dir]
+            + ["-e", args.main]
+        )
+
+        os.replace(tmp_output, args.output)
+
+
 def build_zipapp(args, path_map):
     """Create a self executing python binary using Python 3's built-in
     zipapp module.
@@ -262,6 +293,7 @@ def check_main_module(args, path_map):
 
 
 BUILD_TYPES = {
+    "pex": build_pex,
     "zipapp": build_zipapp,
     "dir": build_install_dir,
     "lib-install": install_library,
diff --git a/build/fbcode_builder/getdeps.py b/build/fbcode_builder/getdeps.py
index cd3e38b8d..cfce7b6b6 100755
--- a/build/fbcode_builder/getdeps.py
+++ b/build/fbcode_builder/getdeps.py
@@ -432,29 +432,35 @@ class InstallSysDepsCmd(ProjectCmdBase):
                 merged += v
                 all_packages[k] = merged
 
-        cmd_args = None
+        cmd_argss = []
         if manager == "rpm":
             packages = sorted(set(all_packages["rpm"]))
             if packages:
-                cmd_args = ["sudo", "dnf", "install", "-y", "--skip-broken"] + packages
+                cmd_argss.append(
+                    ["sudo", "dnf", "install", "-y", "--skip-broken"] + packages
+                )
         elif manager == "deb":
             packages = sorted(set(all_packages["deb"]))
             if packages:
-                cmd_args = [
-                    "sudo",
-                    "--preserve-env=http_proxy",
-                    "apt-get",
-                    "install",
-                    "-y",
-                ] + packages
+                cmd_argss.append(
+                    [
+                        "sudo",
+                        "--preserve-env=http_proxy",
+                        "apt-get",
+                        "install",
+                        "-y",
+                    ]
+                    + packages
+                )
+                cmd_argss.append(["pip", "install", "pex"])
         elif manager == "homebrew":
             packages = sorted(set(all_packages["homebrew"]))
             if packages:
-                cmd_args = ["brew", "install"] + packages
+                cmd_argss.append(["brew", "install"] + packages)
         elif manager == "pacman-package":
             packages = sorted(list(set(all_packages["pacman-package"])))
             if packages:
-                cmd_args = ["pacman", "-S"] + packages
+                cmd_argss.append(["pacman", "-S"] + packages)
         else:
             host_tuple = loader.build_opts.host_type.as_tuple_string()
             print(
@@ -462,7 +468,7 @@ class InstallSysDepsCmd(ProjectCmdBase):
             )
             return
 
-        if cmd_args:
+        for cmd_args in cmd_argss:
             if args.dry_run:
                 print(" ".join(cmd_args))
             else: