numcodecs_wasm/
lib.rs

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