[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

@@ -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()

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