midas/
lib.rs

1#[macro_use]
2extern crate cpython;
3
4use std::{
5    cell::RefCell,
6    collections::hash_map::DefaultHasher,
7    hash::{Hash, Hasher},
8};
9
10use cpython::{PyResult, Python};
11use midas_rs::{default, Float, Int, MidasR as MidasR_, MidasRParams};
12
13py_module_initializer!(midas, initmidas, PyInit_midas, |py, module| {
14    module.add(py, "__doc__", "MidasR implementation")?;
15    module.add(py, "hash", py_fn!(py, hash(string: String)))?;
16    module.add_class::<MidasR>(py)?;
17    Ok(())
18});
19
20fn hash(_py: Python, string: String) -> PyResult<Int> {
21    let mut hasher = DefaultHasher::new();
22    string.hash(&mut hasher);
23    Ok(hasher.finish() as Int)
24}
25
26py_class!(class MidasR |py| {
27    data value: RefCell<MidasR_>;
28
29    def __new__(
30        _cls,
31        rows: Int = default::NUM_ROWS,
32        buckets: Int = default::NUM_BUCKETS,
33        m_value: Int = default::M_VALUE,
34        alpha: Float = default::ALPHA
35    ) -> PyResult<MidasR> {
36        MidasR::create_instance(py, RefCell::new(MidasR_::new(MidasRParams { 
37            rows, buckets, m_value, alpha
38        })))
39    }
40
41    def insert(&self, source: Int, dest: Int, time: Int) -> PyResult<Float> {
42        Ok(self.value(py).borrow_mut().insert((source, dest, time)))
43    }
44
45    def query(&self, source: Int, dest: Int) -> PyResult<Float> {
46        Ok(self.value(py).borrow().query(source, dest))
47    }
48
49    def current_time(&self) -> PyResult<Int> {
50        Ok(self.value(py).borrow().current_time())
51    }
52
53    def alpha(&self) -> PyResult<Float> {
54        Ok(self.value(py).borrow().alpha())
55    }
56});