scouter_types/queue/
types.rs

1use crate::ProfileFuncs;
2use pyo3::prelude::*;
3
4use crate::error::TypeError;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::fmt;
8use std::fmt::Display;
9use std::fmt::Formatter;
10
11#[pyclass]
12#[derive(Clone, Debug, Serialize)]
13pub enum EntityType {
14    Feature,
15    Metric,
16}
17
18#[pyclass]
19#[derive(Clone, Serialize, Deserialize, Debug)]
20pub struct IntFeature {
21    pub name: String,
22    pub value: i64,
23}
24
25#[pymethods]
26impl IntFeature {
27    pub fn __str__(&self) -> String {
28        ProfileFuncs::__str__(self)
29    }
30}
31
32impl IntFeature {
33    pub fn to_float(&self) -> f64 {
34        self.value as f64
35    }
36}
37
38#[pyclass]
39#[derive(Clone, Serialize, Deserialize, Debug)]
40pub struct FloatFeature {
41    pub name: String,
42    pub value: f64,
43}
44
45#[pymethods]
46impl FloatFeature {
47    pub fn __str__(&self) -> String {
48        ProfileFuncs::__str__(self)
49    }
50}
51
52#[pyclass]
53#[derive(Clone, Serialize, Deserialize, Debug)]
54pub struct StringFeature {
55    pub name: String,
56    pub value: String,
57}
58
59#[pymethods]
60impl StringFeature {
61    pub fn __str__(&self) -> String {
62        ProfileFuncs::__str__(self)
63    }
64}
65
66impl StringFeature {
67    pub fn to_float(&self, feature_map: &FeatureMap) -> Result<f64, TypeError> {
68        feature_map
69            .features
70            .get(&self.name)
71            .and_then(|feat_map| {
72                feat_map
73                    .get(&self.value)
74                    .or_else(|| feat_map.get("missing"))
75            })
76            .map(|&val| val as f64)
77            .ok_or(TypeError::MissingStringValueError)
78    }
79}
80
81#[pyclass]
82#[derive(Clone, Serialize, Deserialize, Debug)]
83pub enum Feature {
84    Int(IntFeature),
85    Float(FloatFeature),
86    String(StringFeature),
87}
88
89#[pymethods]
90impl Feature {
91    #[staticmethod]
92    pub fn int(name: String, value: i64) -> Self {
93        Feature::Int(IntFeature { name, value })
94    }
95
96    #[staticmethod]
97    pub fn float(name: String, value: f64) -> Self {
98        Feature::Float(FloatFeature { name, value })
99    }
100
101    #[staticmethod]
102    pub fn string(name: String, value: String) -> Self {
103        Feature::String(StringFeature { name, value })
104    }
105
106    pub fn __str__(&self) -> String {
107        ProfileFuncs::__str__(self)
108    }
109}
110
111impl Feature {
112    pub fn to_float(&self, feature_map: &FeatureMap) -> Result<f64, TypeError> {
113        match self {
114            Feature::Int(feature) => Ok(feature.to_float()),
115            Feature::Float(feature) => Ok(feature.value),
116            Feature::String(feature) => feature.to_float(feature_map),
117        }
118    }
119
120    pub fn name(&self) -> &str {
121        match self {
122            Feature::Int(feature) => &feature.name,
123            Feature::Float(feature) => &feature.name,
124            Feature::String(feature) => &feature.name,
125        }
126    }
127}
128
129impl Display for Feature {
130    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
131        match self {
132            Feature::Int(feature) => write!(f, "{}", feature.value),
133            Feature::Float(feature) => write!(f, "{}", feature.value),
134            Feature::String(feature) => write!(f, "{}", feature.value),
135        }
136    }
137}
138
139#[pyclass]
140#[derive(Clone, Debug, Serialize)]
141pub struct Features {
142    #[pyo3(get)]
143    pub features: Vec<Feature>,
144
145    #[pyo3(get)]
146    pub entity_type: EntityType,
147}
148
149#[pymethods]
150impl Features {
151    #[new]
152    pub fn new(features: Vec<Feature>) -> Self {
153        Features {
154            features,
155            entity_type: EntityType::Feature,
156        }
157    }
158
159    pub fn __str__(&self) -> String {
160        ProfileFuncs::__str__(self)
161    }
162}
163
164impl Features {
165    pub fn iter(&self) -> std::slice::Iter<'_, Feature> {
166        self.features.iter()
167    }
168}
169
170#[pyclass]
171#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)]
172pub struct FeatureMap {
173    #[pyo3(get)]
174    pub features: HashMap<String, HashMap<String, usize>>,
175}
176
177#[pymethods]
178impl FeatureMap {
179    pub fn __str__(&self) -> String {
180        // serialize the struct to a string
181        ProfileFuncs::__str__(self)
182    }
183}
184
185#[pyclass]
186#[derive(Clone, Serialize, Debug)]
187pub struct Metric {
188    pub name: String,
189    pub value: f64,
190}
191
192#[pymethods]
193impl Metric {
194    #[new]
195    pub fn new(name: String, value: f64) -> Self {
196        Metric { name, value }
197    }
198    pub fn __str__(&self) -> String {
199        ProfileFuncs::__str__(self)
200    }
201}
202
203#[pyclass]
204#[derive(Clone, Serialize, Debug)]
205pub struct Metrics {
206    #[pyo3(get)]
207    pub metrics: Vec<Metric>,
208
209    #[pyo3(get)]
210    pub entity_type: EntityType,
211}
212
213#[pymethods]
214impl Metrics {
215    #[new]
216    pub fn new(metrics: Vec<Metric>) -> Self {
217        Metrics {
218            metrics,
219            entity_type: EntityType::Metric,
220        }
221    }
222    pub fn __str__(&self) -> String {
223        ProfileFuncs::__str__(self)
224    }
225}
226
227impl Metrics {
228    pub fn iter(&self) -> std::slice::Iter<'_, Metric> {
229        self.metrics.iter()
230    }
231}
232
233#[derive(Debug)]
234pub enum QueueItem {
235    Features(Features),
236    Metrics(Metrics),
237}
238
239impl QueueItem {
240    /// Helper for extracting an Entity from a Python object
241    pub fn from_py_entity(entity: &Bound<'_, PyAny>) -> Result<Self, TypeError> {
242        let entity_type = entity.getattr("entity_type")?.extract::<EntityType>()?;
243
244        match entity_type {
245            EntityType::Feature => {
246                let features = entity.extract::<Features>()?;
247                Ok(QueueItem::Features(features))
248            }
249            EntityType::Metric => {
250                let metrics = entity.extract::<Metrics>()?;
251                Ok(QueueItem::Metrics(metrics))
252            }
253        }
254    }
255}
256
257pub trait QueueExt: Send + Sync {
258    fn metrics(&self) -> &Vec<Metric>;
259    fn features(&self) -> &Vec<Feature>;
260}
261
262impl QueueExt for Features {
263    fn metrics(&self) -> &Vec<Metric> {
264        // this is not a real implementation, just a placeholder
265        // to satisfy the trait bound
266        static EMPTY: Vec<Metric> = Vec::new();
267        &EMPTY
268    }
269
270    fn features(&self) -> &Vec<Feature> {
271        &self.features
272    }
273}
274
275impl QueueExt for Metrics {
276    fn metrics(&self) -> &Vec<Metric> {
277        &self.metrics
278    }
279
280    fn features(&self) -> &Vec<Feature> {
281        // this is not a real implementation, just a placeholder
282        // to satisfy the trait bound
283        static EMPTY: Vec<Feature> = Vec::new();
284        &EMPTY
285    }
286}
scouter_types/queue/types.rs

scouter_types/queue/
types.rs
