From 0ff14b776f2bfd7672f36d7d082e5c0ba9e57bfe Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Wed, 18 Sep 2024 07:26:26 -0700
Subject: [PATCH] [libshortfin] Fix python stub generation.

This brings pylance back to life for the native bindings. It was giving confusing error squiggles about not being able to use types that were variables, but the root cause was that the backing stubs were not generated properly for all of the module.
---
 libshortfin/python/CMakeLists.txt        | 37 ++++++++++++++++++++----
 libshortfin/python/_shortfin/__init__.py | 35 ++++++++++++----------
 2 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/libshortfin/python/CMakeLists.txt b/libshortfin/python/CMakeLists.txt
index f0048d6a7..f651484f9 100644
--- a/libshortfin/python/CMakeLists.txt
+++ b/libshortfin/python/CMakeLists.txt
@@ -35,18 +35,43 @@ target_link_libraries(shortfin_python_extension
   PRIVATE ${SHORTFIN_LINK_LIBRARY_NAME}
 )
 
-if (SHORTFIN_ENABLE_TRACING)
+function(shortfin_python_stubs build_type)
   nanobind_add_stub(
     shortfin_python_extension_stub
-    MODULE _shortfin_tracy.lib
-    OUTPUT _shortfin_tracy/lib.pyi
+    MODULE _shortfin_${build_type}.lib
+    OUTPUT _shortfin_${build_type}/lib.pyi
     DEPENDS shortfin_python_extension
   )
-else()
+
+endfunction()
+
+function(shortfin_python_stubs build_variant output_root)
+  file(MAKE_DIRECTORY ${output_root})
   nanobind_add_stub(
-    shortfin_python_extension_stub
+    shortfin_python_extension_stub_lib_${build_variant}
     MODULE _shortfin_default.lib
-    OUTPUT _shortfin_default/lib.pyi
+    OUTPUT ${output_root}/lib/__init__.pyi
     DEPENDS shortfin_python_extension
   )
+
+  nanobind_add_stub(
+    shortfin_python_extension_stub_array_${build_variant}
+    MODULE _shortfin_default.lib.array
+    OUTPUT ${output_root}/lib/array.pyi
+    DEPENDS shortfin_python_extension
+  )
+
+  nanobind_add_stub(
+    shortfin_python_extension_stub_local_${build_variant}
+    MODULE _shortfin_default.lib.local
+    OUTPUT ${output_root}/lib/local.pyi
+    DEPENDS shortfin_python_extension
+  )
+endfunction()
+
+# Generate the same stubs against the default build for each variant but
+# output the files to the right package.
+shortfin_python_stubs(default ${CMAKE_CURRENT_BINARY_DIR}/_shortfin_default)
+if (SHORTFIN_ENABLE_TRACING)
+  shortfin_python_stubs(tracy ${CMAKE_CURRENT_BINARY_DIR}/_shortfin_tracy)
 endif()
diff --git a/libshortfin/python/_shortfin/__init__.py b/libshortfin/python/_shortfin/__init__.py
index 79fbd924a..9bfa3a497 100644
--- a/libshortfin/python/_shortfin/__init__.py
+++ b/libshortfin/python/_shortfin/__init__.py
@@ -7,23 +7,28 @@
 # The proper way to import this package is via:
 #   from _shortfin import lib as sfl
 
+from typing import TYPE_CHECKING
+
 import os
 import sys
 import warnings
 
-variant = os.getenv("SHORTFIN_PY_RUNTIME", "default")
-
-if variant == "tracy":
-    try:
-        from _shortfin_tracy import lib
-    except ModuleNotFoundError as e:
-        raise ModuleNotFoundError(
-            "Shortfin Tracy runtime requested via SHORTFIN_PY_RUNTIME but it is not enabled in this build"
-        )
-    print("-- Using Tracy runtime (SHORTFIN_PY_RUNTIME=tracy)", file=sys.stderr)
-else:
-    if variant != "default":
-        warnings.warn(
-            f"Unknown value for SHORTFIN_PY_RUNTIME env var ({variant}): Using default"
-        )
+if TYPE_CHECKING:
     from _shortfin_default import lib
+else:
+    variant = os.getenv("SHORTFIN_PY_RUNTIME", "default")
+
+    if variant == "tracy":
+        try:
+            from _shortfin_tracy import lib
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError(
+                "Shortfin Tracy runtime requested via SHORTFIN_PY_RUNTIME but it is not enabled in this build"
+            )
+        print("-- Using Tracy runtime (SHORTFIN_PY_RUNTIME=tracy)", file=sys.stderr)
+    else:
+        if variant != "default":
+            warnings.warn(
+                f"Unknown value for SHORTFIN_PY_RUNTIME env var ({variant}): Using default"
+            )
+        from _shortfin_default import lib