numcodecs_wasm/
lib.rs

1use numcodecs_python::{PyCodec, PyCodecAdapter, PyCodecClass};
2use numcodecs_wasm_host_reproducible::{ReproducibleWasmCodec, ReproducibleWasmCodecType};
3use pyo3::{exceptions::PyTypeError, prelude::*};
4
5// FIXME: https://github.com/bytecodealliance/rustix/issues/1620
6use ::memfd as _;
7
8mod engine;
9
10use engine::{Engine, default_engine};
11
12#[pymodule]
13#[pyo3(name = "_wasm")]
14fn wasm<'py>(py: Python<'py>, module: &Bound<'py, PyModule>) -> Result<(), PyErr> {
15    let logger = pyo3_log::Logger::new(py, pyo3_log::Caching::Nothing)?;
16    logger
17        .install()
18        .map_err(|err| pyo3_error::PyErrChain::new(py, err))?;
19
20    module.add_function(wrap_pyfunction!(create_codec_class, module)?)?;
21    module.add_function(wrap_pyfunction!(read_codec_instruction_counter, module)?)?;
22
23    Ok(())
24}
25
26#[pyfunction]
27#[pyo3(name = "_create_codec_class")]
28fn create_codec_class<'py>(
29    py: Python<'py>,
30    module: &Bound<'py, PyModule>,
31    wasm: Vec<u8>,
32) -> Result<Bound<'py, PyCodecClass>, PyErr> {
33    let engine = default_engine(py)?;
34
35    // TODO: we should allow restricting the codecs that the reproducible codec can 'see'
36    let codec_ty = ReproducibleWasmCodecType::new_with_registry(
37        engine,
38        wasm,
39        numcodecs_python::PyCodecRegistryHandle,
40    )
41    .map_err(|err| pyo3_error::PyErrChain::new(py, err))?;
42
43    let codec_class = numcodecs_python::export_codec_class(py, codec_ty, module.as_borrowed())?;
44
45    Ok(codec_class)
46}
47
48#[pyfunction]
49#[pyo3(name = "_read_codec_instruction_counter")]
50fn read_codec_instruction_counter<'py>(
51    py: Python<'py>,
52    codec: &Bound<'py, PyCodec>,
53) -> Result<u64, PyErr> {
54    let Some(instruction_counter) =
55        PyCodecAdapter::with_downcast(py, codec, |codec: &ReproducibleWasmCodec<Engine>| {
56            codec.instruction_counter()
57        })
58        .transpose()
59        .map_err(|err| pyo3_error::PyErrChain::new(py, err))?
60    else {
61        return Err(PyTypeError::new_err(
62            "`codec` is not a wasm codec, only wasm codecs have instruction counts",
63        ));
64    };
65
66    Ok(instruction_counter.0)
67}
numcodecs_wasm/lib.rs

numcodecs_wasm/
lib.rs
