use std::prelude::v1::*;
use pyo3::prelude::*;
use crate::{
pybindings::{PyReadonlyFloatArray, PyReadonlyFloatArray1},
symbol::huffman,
};
#[pymodule]
#[pyo3(name = "huffman")]
pub fn init_module(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_class::<EncoderHuffmanTree>()?;
module.add_class::<DecoderHuffmanTree>()?;
Ok(())
}
#[pyclass]
#[derive(Debug)]
pub struct EncoderHuffmanTree {
pub(crate) inner: huffman::EncoderHuffmanTree,
}
#[pymethods]
impl EncoderHuffmanTree {
#[new]
#[pyo3(signature = (probabilities))]
pub fn new(probabilities: PyReadonlyFloatArray1<'_>) -> PyResult<Self> {
let inner = match probabilities {
PyReadonlyFloatArray::F32(probabilities) => {
huffman::EncoderHuffmanTree::from_float_probabilities::<f32, _>(
probabilities.as_array(),
)
}
PyReadonlyFloatArray::F64(probabilities) => {
huffman::EncoderHuffmanTree::from_float_probabilities::<f64, _>(
probabilities.as_array(),
)
}
}?;
Ok(Self { inner })
}
}
#[pyclass]
#[derive(Debug)]
pub struct DecoderHuffmanTree {
pub(crate) inner: huffman::DecoderHuffmanTree,
}
#[pymethods]
impl DecoderHuffmanTree {
#[new]
#[pyo3(signature = (probabilities))]
pub fn new(probabilities: PyReadonlyFloatArray1<'_>) -> PyResult<Self> {
let inner = match probabilities {
PyReadonlyFloatArray::F32(probabilities) => {
huffman::DecoderHuffmanTree::from_float_probabilities::<f32, _>(
probabilities.as_array(),
)
}
PyReadonlyFloatArray::F64(probabilities) => {
huffman::DecoderHuffmanTree::from_float_probabilities::<f64, _>(
probabilities.as_array(),
)
}
}?;
Ok(Self { inner })
}
}