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}