From fe2a559d93e57ded681a2025f4b560541c6f988b Mon Sep 17 00:00:00 2001
From: Aaron Miller <78561124+aaron-skydio@users.noreply.github.com>
Date: Mon, 5 Jun 2023 23:45:00 +0000
Subject: [PATCH] [SymForce-External] Fix symengine_wrapper search

Fix symengine_wrapper search

So that it works on different platforms that e.g. use `local/lib` or
`site-packages`.

We were doing this correctly for editable installs but not non-editable
installs

Longer-term I'm leaning more and more towards requiring building
our copy of symengine separately first, but this is an easy strict
improvement for now

Fixes symforce-org/symforce#330

Topic: sf-wrapper-glob
Reviewers: bradley,nathan,chao,ryan-b
Closes symforce-org/symforce#342
GitOrigin-RevId: fce960b93a984811754da248fed658641fa16367
---
 setup.py | 39 ++++++++++++++++++++++++++-------------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/setup.py b/setup.py
index 19f7978ea..caf1fa505 100644
--- a/setup.py
+++ b/setup.py
@@ -9,6 +9,7 @@
 import re
 import subprocess
 import sys
+import typing as T
 from pathlib import Path
 
 from setuptools import Extension
@@ -145,11 +146,8 @@ def run(self) -> None:
             # in the symenginepy/symengine source directory. Everything is already there
             # except the compiled symengine_wrapper extension module, so we copy that there
             # as well.
-            symengine_wrapper = next(
-                build_temp_path.glob(
-                    f"symengine_install/**/lib/python{sys.version_info.major}.{sys.version_info.minor}/*-packages/symengine/lib/{self.get_ext_filename('symengine_wrapper')}"
-                ),
-                None,
+            symengine_wrapper = maybe_find_symengine_wrapper(
+                build_temp_path, self.get_ext_filename("symengine_wrapper")
             )
             # NOTE(brad) For some reason that I don't fully understand, the symengine_wrapper
             # shared library is neither present in symengine_install nor needed in the source
@@ -264,6 +262,28 @@ def filter_local(s: str) -> str:
         super().run()
 
 
+def maybe_find_symengine_wrapper(build_dir: Path, ext_filename: str) -> T.Optional[Path]:
+    symengine_wrapper_candidates = list(
+        build_dir.glob(
+            f"symengine_install/**/lib/python{sys.version_info.major}.{sys.version_info.minor}/*-packages/symengine/lib/{ext_filename}"
+        )
+    )
+
+    if len(symengine_wrapper_candidates) > 1:
+        raise FileNotFoundError(
+            f"Expected to find exactly one symengine_wrapper.so, but found {len(symengine_wrapper_candidates)}: {symengine_wrapper_candidates}"
+        )
+
+    return next(iter(symengine_wrapper_candidates), None)
+
+
+def find_symengine_wrapper(build_dir: Path, ext_filename: str) -> Path:
+    symengine_wrapper = maybe_find_symengine_wrapper(build_dir, ext_filename)
+    if symengine_wrapper is None:
+        raise FileNotFoundError(f"Could not find symengine_wrapper.so in {build_dir}")
+    return symengine_wrapper
+
+
 class InstallWithExtras(install):
     """
     Custom install step that:
@@ -284,14 +304,7 @@ def run(self) -> None:
         # include the additional Cython sources that symenginepy includes in their distributions,
         # but I'm honestly not sure why they include them or why you'd need them.
         self.copy_file(
-            build_dir
-            / "symengine_install"
-            / "lib"
-            / f"python{sys.version_info.major}.{sys.version_info.minor}"
-            / "site-packages"
-            / "symengine"
-            / "lib"
-            / build_ext_obj.get_ext_filename("symengine_wrapper"),
+            find_symengine_wrapper(build_dir, build_ext_obj.get_ext_filename("symengine_wrapper")),
             Path.cwd()
             / self.install_platlib  # type: ignore[attr-defined]
             / "symengine"