Support building local codecs with numcodecs-wasm-materialize

Author: juntyr

.github/workflows/ci.yml

@@ -296,7 +296,7 @@ jobs:
           rm .python-version
           uv python install ${{ matrix.python }}
           uv sync && uv pip install .
-          uv run python3 -m numcodecs_wasm_materialize zfp
+          uv run python3 -m numcodecs_wasm_materialize zfp --local
 
       - name: Test importing the ZFP WASM codec
         run: |

crates/numcodecs-wasm-builder/src/main.rs

@@ -3,7 +3,7 @@
 
 use std::{
     collections::HashMap,
-    fs, io,
+    env, fs, io,
     path::{Path, PathBuf},
     process::Command,
     str::FromStr,
@@ -30,6 +30,10 @@ struct Args {
     /// Path to which the wasm file is output
     #[arg(long, short)]
     output: PathBuf,
+
+    /// Compile the local crate instead of the published one
+    #[arg(long)]
+    local: bool,
 }
 
 fn main() -> io::Result<()> {
@@ -47,8 +51,13 @@ fn main() -> io::Result<()> {
     eprintln!("creating {}", target_dir.display());
     fs::create_dir_all(&target_dir)?;
 
-    let crate_dir =
-        create_codec_wasm_component_crate(&scratch_dir, &args.crate_, &args.version, &args.codec)?;
+    let crate_dir = create_codec_wasm_component_crate(
+        &scratch_dir,
+        &args.crate_,
+        &args.version,
+        &args.codec,
+        args.local,
+    )?;
     copy_buildenv_to_crate(&crate_dir)?;
 
     let nix_env = NixEnv::new(&crate_dir)?;
@@ -72,6 +81,7 @@ fn create_codec_wasm_component_crate(
     crate_: &str,
     version: &Version,
     codec: &str,
+    local: bool,
 ) -> io::Result<PathBuf> {
     let crate_dir = scratch_dir.join(format!("{crate_}-wasm-{version}"));
     eprintln!("crate_dir={}", crate_dir.display());
@@ -81,6 +91,25 @@ fn create_codec_wasm_component_crate(
     }
     fs::create_dir_all(&crate_dir)?;
 
+    let (numcodecs_wasm_logging_path, numcodecs_wasm_guest_path, numcodecs_my_codec_path) = if local
+    {
+        let numcodecs = Path::new(env!("CARGO_MANIFEST_DIR")).join("..").join("..");
+        eprintln!("looking for local workspace in {}", numcodecs.display());
+        let numcodecs = numcodecs.canonicalize()?;
+        let numcodecs_wasm_logging = numcodecs.join("crates").join("numcodecs-wasm-logging");
+        let numcodecs_wasm_guest = numcodecs.join("crates").join("numcodecs-wasm-guest");
+        let numcodecs_my_codec = numcodecs
+            .join("codecs")
+            .join(crate_.strip_prefix("numcodecs-").unwrap_or(crate_));
+        (
+            format!(r#" path = "{}","#, numcodecs_wasm_logging.display()),
+            format!(r#" path = "{}","#, numcodecs_wasm_guest.display()),
+            format!(r#" path = "{}","#, numcodecs_my_codec.display()),
+        )
+    } else {
+        (String::new(), String::new(), String::new())
+    };
+
     fs::write(
         crate_dir.join("Cargo.toml"),
         format!(
@@ -95,9 +124,9 @@ edition = "2024"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-numcodecs-wasm-logging = {{ version = "0.2", default-features = false }}
-numcodecs-wasm-guest = {{ version = "0.3", default-features = false }}
-numcodecs-my-codec = {{ package = "{crate_}", version = "{version}", default-features = false }}
+numcodecs-wasm-logging = {{ version = "0.2",{numcodecs_wasm_logging_path} default-features = false }}
+numcodecs-wasm-guest = {{ version = "0.3",{numcodecs_wasm_guest_path} default-features = false }}
+numcodecs-my-codec = {{ package = "{crate_}", version = "{version}",{numcodecs_my_codec_path} default-features = false }}
     "#
         ),
     )?;

py/numcodecs-wasm-materialize/src/numcodecs_wasm_materialize/__main__.py

@@ -1,8 +1,8 @@
+import argparse
 import re
 import shlex
 import shutil
 import subprocess
-import sys
 from pathlib import Path
 
 import toml
@@ -13,12 +13,23 @@
 template_path = repo_path / "py" / "numcodecs-wasm-template"
 dist_path = repo_path / "py" / "numcodecs-wasm-materialize" / "dist"
 
-codec_pattern = sys.argv[1]
+parser = argparse.ArgumentParser(prog="numcodecs_wasm_materialize")
+parser.add_argument(
+    "codec",
+    type=re.compile,
+    help="regex pattern for the codecs to materialize",
+)
+parser.add_argument(
+    "--local",
+    action="store_true",
+    help="compile the local codec instead of the published one",
+)
+args = parser.parse_args()
 
 for c in (repo_path / "codecs").iterdir():
     codec = c.name
 
-    if not codec_pattern or re.match(codec_pattern, codec) is None:
+    if not args.codec or re.match(args.codec, codec) is None:
         continue
 
     codec_crate_path = repo_path / "codecs" / codec
@@ -77,10 +88,11 @@
     subprocess.run(
         shlex.split(
             "cargo run -p numcodecs-wasm-builder --"
-            f" --crate numcodecs-{templates['crate-suffix']}"
-            f" --version {templates['crate-version']}"
-            f" --codec {templates['codec-path']}"
-            f" --output {staging_path / 'src' / ('numcodecs_wasm_' + templates['package_suffix']) / 'codec.wasm'}"
+            + f" --crate numcodecs-{templates['crate-suffix']}"
+            + f" --version {templates['crate-version']}"
+            + f" --codec {templates['codec-path']}"
+            + f" --output {staging_path / 'src' / ('numcodecs_wasm_' + templates['package_suffix']) / 'codec.wasm'}"
+            + (" --local" if args.local else "")
         ),
         check=True,
     )
@@ -98,16 +110,16 @@
     subprocess.run(
         shlex.split(
             f"uv run python3 {Path(__file__).parent / 'stub.py'} "
-            f"{'numcodecs_wasm_' + templates['package_suffix']} src"
+            + f"{'numcodecs_wasm_' + templates['package_suffix']} src"
         ),
         check=True,
         cwd=staging_path,
     )
     subprocess.run(
         shlex.split(
             f'uv run python3 -c "from {"numcodecs_wasm_" + templates["package_suffix"]} '
-            f'import {templates["CodecName"]} as Codec; '
-            f'assert Codec.codec_id == {templates["codec-id"]!r}, Codec.codec_id"'
+            + f"import {templates['CodecName']} as Codec; "
+            + f'assert Codec.codec_id == {templates["codec-id"]!r}, Codec.codec_id"'
         ),
         check=True,
         cwd=staging_path,