Merge pull request #76 from trungleduc/conda-format

Support direct download for no-arch packages
emscripten-forge · Oct 22, 2024 · 28d35bd · 28d35bd
2 parents b27943b + b7dcff8
commit 28d35bd
Show file tree

Hide file tree

Showing 10 changed files with 310 additions and 18 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -53,7 +53,7 @@ jobs:
             -c https://repo.mamba.pm/emscripten-forge \
             -c https://repo.mamba.pm/conda-forge \
             --yes \
-            python pybind11 nlohmann_json pybind11_json numpy "pytest==7.1.1" bzip2 sqlite zlib libffi exceptiongroup
+            python pybind11 nlohmann_json pybind11_json numpy "pytest==7.1.1" bzip2 sqlite zlib zstd libffi exceptiongroup
 
 
           mkdir build

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -83,6 +83,7 @@ set(PYJS_HEADERS
     include/pyjs/export_py_object.hpp
     include/pyjs/export_pyjs_module.hpp
     include/pyjs/untar.hpp
+    include/pyjs/install_conda_file.hpp
     include/pyjs/inflate.hpp
     ${CMAKE_CURRENT_BINARY_DIR}/pyjs_pre.js
     ${CMAKE_CURRENT_BINARY_DIR}/pyjs_post.js
@@ -110,6 +111,7 @@ add_library(pyjs STATIC
     src/js_timestamp.cpp
     src/inflate.cpp
     src/untar.cpp
+    src/install_conda_file.cpp
     ${PYCPPSOURCES}
 )
 
@@ -153,6 +155,7 @@ SET(PYTHON_UTIL_LIBS
     ${CMAKE_INSTALL_PREFIX}/lib/libz.a
     ${CMAKE_INSTALL_PREFIX}/lib/libsqlite3.a
     ${CMAKE_INSTALL_PREFIX}/lib/libffi.a
+    ${CMAKE_INSTALL_PREFIX}/lib/libzstd.a
 )
 
 

diff --git a/build_mkdocs.sh b/build_mkdocs.sh
@@ -28,7 +28,7 @@ if [ ! -d "$WASM_ENV_PREFIX" ]; then
             -c https://repo.mamba.pm/conda-forge \
             --yes \
             python pybind11 nlohmann_json pybind11_json numpy \
-            bzip2 sqlite zlib libffi exceptiongroup \
+            bzip2 sqlite zlib zstd libffi exceptiongroup \
             "xeus<4" "xeus-lite<2" xeus-python "xeus-javascript>=0.3.2" xtl "ipython=8.22.2=py311had7285e_1" "traitlets>=5.14.2"
 
 else

diff --git a/include/pyjs/install_conda_file.hpp b/include/pyjs/install_conda_file.hpp
@@ -0,0 +1,12 @@
+#pragma once
+
+#include <emscripten/bind.h>
+
+namespace pyjs
+{
+
+    em::val install_conda_file(const std::string& zstd_file_path,
+                               const std::string& working_dir,
+                               const std::string& path);
+
+}
diff --git a/include/pyjs/pre_js/load_pkg.js b/include/pyjs/pre_js/load_pkg.js
@@ -67,6 +67,69 @@ def _py_untar(tarball_path, target_dir):
 }
 
 
+Module["_unzip_from_python"] = function(tarball_path, target_dir) {
+    Module.exec(`
+def _py_unzip(tarball_path, target_dir):
+    import json
+    from pathlib import Path
+    import zipfile
+    
+    target = Path(target_dir)
+    target.mkdir(parents=True, exist_ok=True)
+    pkg_file = {"name": "", "path": ""}
+    with zipfile.ZipFile(tarball_path, mode="r") as archive:
+        
+        for filename in archive.namelist():
+            if filename.startswith("pkg-"):
+                pkg_file["name"] = filename
+                pkg_file["path"] = str(target / filename)
+                archive.extract(filename, target_dir)    
+                break
+    return json.dumps(pkg_file)
+
+`)
+    let extracted_file = Module.eval(`_py_unzip("${tarball_path}", "${target_dir}")`)
+
+    return JSON.parse(extracted_file)
+}
+
+Module["_install_conda_file_from_python"] = function(tarball_path, target_dir) {
+    Module.exec(`
+def _py_unbz2(tarball_path, target_dir):
+    import json
+    from pathlib import Path
+    import tarfile
+    import shutil
+    import os
+    import sys
+    
+    target = Path(target_dir)
+    prefix = Path(sys.prefix)
+    try:
+        with tarfile.open(tarball_path) as tar:
+            tar.extractall(target_dir)
+
+        src = target / "site-packages"
+        dest = prefix / "lib/python3.11/site-packages"
+        shutil.copytree(src, dest, dirs_exist_ok=True)
+        for folder in ["etc", "share"]:
+            src = target / folder
+            dest = prefix / folder
+            if src.exists():
+                shutil.copytree(src, dest, dirs_exist_ok=True)
+        shutil.rmtree(target)
+    except Exception as e:
+        print("ERROR",e)
+        raise e
+    
+    return json.dumps([])
+
+`)
+    let extracted_file = Module.eval(`_py_unbz2("${tarball_path}", "${target_dir}")`)
+
+    return JSON.parse(extracted_file)
+}
+
 
 
 
@@ -108,20 +171,55 @@ Module["bootstrap_from_empack_packed_environment"] = async function
                 pkg,
                 verbose
             ) {
-            const package_url = pkg?.url ?? `${package_tarballs_root_url}/${pkg.filename}`;
-            if (verbose) {
-                console.log(`!!fetching pkg ${pkg.name} from ${package_url}`)
-            }
-            let byte_array = await fetchByteArray(package_url)
-            const tarball_path = `/package_tarballs/${pkg.filename}`;
-            Module.FS.writeFile(tarball_path, byte_array);
-            if(verbose){
-                console.log(`!!extract ${tarball_path} (${byte_array.length} bytes)`)
+              const package_url =
+                pkg?.url ?? `${package_tarballs_root_url}/${pkg.filename}`;
+              if (verbose) {
+                console.log(`!!fetching pkg ${pkg.name} from ${package_url}`);
+              }
+              let byte_array = await fetchByteArray(package_url);
+              const tarball_path = `/package_tarballs/${pkg.filename}`;
+              Module.FS.writeFile(tarball_path, byte_array);
+              if (verbose) {
+                console.log(
+                  `!!extract ${tarball_path} (${byte_array.length} bytes)`
+                );
+              }
+
+              if (verbose) {
+                console.log("await python_is_ready_promise");
+              }
+              await python_is_ready_promise;
+
+              if (package_url.toLowerCase().endsWith(".conda")) {
+                // Conda v2 packages
+                if (verbose) {
+                  console.log(
+                    `!!extract conda package ${package_url} (${byte_array.length} bytes)`
+                  );
+                }
+                const dest = `/conda_packages/${pkg.name}`;
+                const pkg_file = Module["_unzip_from_python"](
+                  tarball_path,
+                  dest
+                );
+                return Module._install_conda_file(pkg_file.path, dest, prefix);
+              } else if (package_url.toLowerCase().endsWith(".tar.bz2")) {
+                // Conda v1 packages
+                if (verbose) {
+                  console.log(
+                    `!!extract conda package ${package_url} (${byte_array.length} bytes)`
+                  );
+                }
+                const dest = `/conda_packages/${pkg.name}`;
+                return Module["_install_conda_file_from_python"](
+                  tarball_path,
+                  dest
+                );
+              } else {
+                // Pre-relocated packages
+                return Module["_untar_from_python"](tarball_path);
+              }
             }
-            if(verbose){console.log("await python_is_ready_promise");}     
-            await python_is_ready_promise;
-            return Module["_untar_from_python"](tarball_path);
-        }
 
 
         async function bootstrap_python(prefix, package_tarballs_root_url, python_package, verbose) {

diff --git a/include/pyjs/untar.hpp b/include/pyjs/untar.hpp
@@ -1,9 +1,10 @@
 #pragma once
 
 #include <emscripten/bind.h>
+#include <emscripten/val.h>
 
+namespace em = emscripten;
 namespace pyjs{
-
     em::val untar(const std::string &tar_path, const std::string &path);
-
+    void untar_impl(FILE *a, const char *path, em::val & shared_libraraies);
 }
diff --git a/src/convert.cpp b/src/convert.cpp
@@ -9,6 +9,7 @@
 namespace py = pybind11;
 namespace em = emscripten;
 
+
 namespace pyjs
 {
     std::pair<em::val,bool> implicit_py_to_js(py::object& py_ret)

diff --git a/src/export_js_module.cpp b/src/export_js_module.cpp
@@ -4,6 +4,7 @@
 #include <pyjs/export_py_object.hpp>
 #include <pyjs/convert.hpp>
 #include <pyjs/untar.hpp>
+#include <pyjs/install_conda_file.hpp>
 
 #include <pybind11/embed.h>
 #include <emscripten/bind.h>
@@ -105,6 +106,7 @@ namespace pyjs
 
 
         em::function("_untar", &untar);
+        em::function("_install_conda_file", &install_conda_file);
         em::function("setenv", &set_env);
 
         // py-object (proxy)