From 9cc9b623bedd8a5c84da5922ef61212e88d8b566 Mon Sep 17 00:00:00 2001
From: Juniper Tyree <50025784+juntyr@users.noreply.github.com>
Date: Sat, 15 Feb 2025 05:13:15 +0000
Subject: [PATCH] Try to fix package imports

---
 tools/clean-dist.py |  4 +++
 tools/makelock.mjs  | 60 +++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/tools/clean-dist.py b/tools/clean-dist.py
index 547d8b33e..65358f83d 100644
--- a/tools/clean-dist.py
+++ b/tools/clean-dist.py
@@ -28,5 +28,9 @@
 for name in packages_to_delete:
     del lock["packages"][name]
 
+for package in lock["packages"].values():
+    package["depends"] = sorted(package["depends"])
+    package["imports"] = sorted(package["imports"])
+
 with lock_path.open("w") as f:
     json.dump(lock, f, sort_keys=True)
diff --git a/tools/makelock.mjs b/tools/makelock.mjs
index 7be317772..a52dce34e 100644
--- a/tools/makelock.mjs
+++ b/tools/makelock.mjs
@@ -10,10 +10,60 @@ const requirements = readFileSync(requirements_path, { encoding: 'utf8' });
 const py = await loadPyodide({ packages: ["micropip"] });
 
 await py.runPythonAsync(`
+import importlib.metadata
 import json
+from pathlib import Path
 
 import micropip
 
+def get_imports_for_package(p: str) -> list[str]:
+    imports = set()
+
+    tree = dict()
+    for f in importlib.metadata.files(p):
+        # ignore special folders
+        if Path(f.parts[0]).suffix in [".libs", ".dist-info", ".data"]:
+            continue
+        
+        # include top-level single-file packages
+        if len(f.parts) == 1 and f.suffix == ".py":
+            imports.add(f.stem)
+            continue
+
+        # build a tree of all other files
+        t = tree
+        for r in f.parts:
+            if t.get(r, None) is None:
+                t[r] = dict()
+            t = t[r]
+
+    # extract folders that only have folders but no files as children,
+    #  these are package candidates
+    queue = [([k], t) for k, t in tree.items() if len(t) > 0]
+    while len(queue) > 0:
+        ps, tree = queue.pop()
+        imports.add('.'.join(ps))
+
+        is_package = True
+
+        add_to_queue = []
+        for k, t in tree.items():
+            if len(t) > 0:
+                add_to_queue.append((ps + [k], t))
+            else:
+                is_package = False
+
+        if is_package:
+            queue += add_to_queue
+
+    # remove prefixes from the list
+    new_imports = []
+    for i in imports:
+        if not any(j.startswith(i) for j in imports if j != i):
+            new_imports.append(i)
+    
+            return new_imports
+
 micropip.set_index_urls([
     "http://0.0.0.0:8000/pypa/simple/{package_name}/",
     "https://pypi.org/pypi/{package_name}/json",
@@ -29,10 +79,16 @@ await micropip.install([
     if len(r) > 0 and not r.startswith('#')
 ], verbose=True)
 
-lock = micropip.freeze().replace("http://0.0.0.0:8000/dist/", "")
+lock = json.loads(
+    micropip.freeze().replace("http://0.0.0.0:8000/dist/", "")
+)
+
+for package in lock["packages"].values():
+    package["depends"] = sorted(package["depends"])
+    package["imports"] = sorted(get_imports_for_package(package["name"]))
 
 with open("/pyodide-lock.json", "w") as f:
-    json.dump(json.loads(lock), f, sort_keys=True)
+    json.dump(lock, f, sort_keys=True)
 `);
 
 const lock = py.FS.readFile("/pyodide-lock.json", { encoding: 'utf8' });