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(py, codec, |codec: &ReproducibleWasmCodec<Engine>| {
48 codec.instruction_counter()
49 })
50 .transpose()
51 .map_err(|err| pyo3_error::PyErrChain::new(py, err))?
52 else {
53 return Err(PyTypeError::new_err(
54 "`codec` is not a wasm codec, only wasm codecs have instruction counts",
55 ));
56 };
57
58 Ok(instruction_counter.0)
59}