Skip to main content

dbn/python/
conversions.rs

1use std::ffi::c_char;
2
3use pyo3::{
4    conversion::IntoPyObjectExt,
5    intern,
6    prelude::*,
7    types::{PyDateTime, PyDict, PyTzInfo},
8};
9
10use crate::{
11    python::PyFieldDesc, BidAskPair, ConsolidatedBidAskPair, HasRType, WithTsOut, UNDEF_TIMESTAMP,
12};
13
14pub fn char_to_c_char(c: char) -> crate::Result<c_char> {
15    if c.is_ascii() {
16        Ok(c as c_char)
17    } else {
18        Err(crate::Error::Conversion {
19            input: c.to_string(),
20            desired_type: "c_char",
21        })
22    }
23}
24
25impl<const N: usize> PyFieldDesc for [BidAskPair; N] {
26    fn field_dtypes(_field_name: &str) -> Vec<(String, String)> {
27        let mut res = Vec::new();
28        let field_dtypes = BidAskPair::field_dtypes("");
29        for level in 0..N {
30            let mut dtypes = field_dtypes.clone();
31            for dtype in dtypes.iter_mut() {
32                dtype.0.push_str(&format!("_{level:02}"));
33            }
34            res.extend(dtypes);
35        }
36        res
37    }
38
39    fn price_fields(_field_name: &str) -> Vec<String> {
40        append_level_suffix::<N>(BidAskPair::price_fields(""))
41    }
42
43    fn ordered_fields(_field_name: &str) -> Vec<String> {
44        append_level_suffix::<N>(BidAskPair::ordered_fields(""))
45    }
46}
47
48impl<const N: usize> PyFieldDesc for [ConsolidatedBidAskPair; N] {
49    fn field_dtypes(_field_name: &str) -> Vec<(String, String)> {
50        let mut res = Vec::new();
51        let field_dtypes = ConsolidatedBidAskPair::field_dtypes("");
52        for level in 0..N {
53            let mut dtypes = field_dtypes.clone();
54            for dtype in dtypes.iter_mut() {
55                dtype.0.push_str(&format!("_{level:02}"));
56            }
57            res.extend(dtypes);
58        }
59        res
60    }
61
62    fn price_fields(_field_name: &str) -> Vec<String> {
63        append_level_suffix::<N>(ConsolidatedBidAskPair::price_fields(""))
64    }
65
66    fn ordered_fields(_field_name: &str) -> Vec<String> {
67        append_level_suffix::<N>(ConsolidatedBidAskPair::ordered_fields(""))
68    }
69
70    fn hidden_fields(_field_name: &str) -> Vec<String> {
71        append_level_suffix::<N>(ConsolidatedBidAskPair::hidden_fields(""))
72    }
73}
74
75pub fn append_level_suffix<const N: usize>(fields: Vec<String>) -> Vec<String> {
76    let mut res = Vec::new();
77    for level in 0..N {
78        let mut fields = fields.clone();
79        for field in fields.iter_mut() {
80            field.push_str(&format!("_{level:02}"));
81        }
82        res.extend(fields);
83    }
84    res
85}
86
87/// `WithTsOut` adds a `ts_out` field to the main record when converted to Python.
88impl<'py, R> IntoPyObject<'py> for WithTsOut<R>
89where
90    R: HasRType + IntoPyObject<'py>,
91{
92    type Target = PyAny;
93    type Output = Bound<'py, PyAny>;
94    type Error = PyErr;
95
96    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
97        let obj = self.rec.into_bound_py_any(py)?;
98        obj.setattr(intern!(py, "ts_out"), self.ts_out).unwrap();
99        Ok(obj)
100    }
101}
102
103pub fn new_py_timestamp_or_datetime(
104    py: Python<'_>,
105    timestamp: u64,
106) -> PyResult<Option<Bound<'_, PyAny>>> {
107    if timestamp == UNDEF_TIMESTAMP {
108        return Ok(None);
109    }
110    if let Ok(pandas) = PyModule::import(py, intern!(py, "pandas")) {
111        let kwargs = PyDict::new(py);
112        if kwargs.set_item(intern!(py, "utc"), true).is_ok()
113            && kwargs
114                .set_item(intern!(py, "errors"), intern!(py, "coerce"))
115                .is_ok()
116            && kwargs
117                .set_item(intern!(py, "unit"), intern!(py, "ns"))
118                .is_ok()
119        {
120            return pandas
121                .call_method(intern!(py, "to_datetime"), (timestamp,), Some(&kwargs))
122                .map(|o| Some(o.into_pyobject(py).unwrap()));
123        }
124    }
125    let utc_tz = PyTzInfo::utc(py)?;
126    let timestamp_ms = timestamp as f64 / 1_000_000.0;
127    PyDateTime::from_timestamp(py, timestamp_ms, Some(&utc_tz))
128        .map(|o| Some(o.into_pyobject(py).unwrap().into_any()))
129}