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    #[staticmethod]
107    pub fn categorical(name: String, value: String) -> Self {
108        Feature::String(StringFeature { name, value })
109    }
110
111    pub fn __str__(&self) -> String {
112        ProfileFuncs::__str__(self)
113    }
114}
115
116impl Feature {
117    pub fn to_float(&self, feature_map: &FeatureMap) -> Result<f64, TypeError> {
118        match self {
119            Feature::Int(feature) => Ok(feature.to_float()),
120            Feature::Float(feature) => Ok(feature.value),
121            Feature::String(feature) => feature.to_float(feature_map),
122        }
123    }
124
125    pub fn name(&self) -> &str {
126        match self {
127            Feature::Int(feature) => &feature.name,
128            Feature::Float(feature) => &feature.name,
129            Feature::String(feature) => &feature.name,
130        }
131    }
132
133    pub fn to_usize(&self, feature_map: &FeatureMap) -> Result<usize, TypeError> {
134        match self {
135            Feature::Int(f) => Ok(f.value as usize),
136            Feature::Float(f) => Ok(f.value as usize),
137            Feature::String(f) => Ok(f.to_float(feature_map)? as usize),
138        }
139    }
140}
141
142impl Display for Feature {
143    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
144        match self {
145            Feature::Int(feature) => write!(f, "{}", feature.value),
146            Feature::Float(feature) => write!(f, "{}", feature.value),
147            Feature::String(feature) => write!(f, "{}", feature.value),
148        }
149    }
150}
151
152#[pyclass]
153#[derive(Clone, Debug, Serialize)]
154pub struct Features {
155    #[pyo3(get)]
156    pub features: Vec<Feature>,
157
158    #[pyo3(get)]
159    pub entity_type: EntityType,
160}
161
162#[pymethods]
163impl Features {
164    #[new]
165    pub fn new(features: Vec<Feature>) -> Self {
166        Features {
167            features,
168            entity_type: EntityType::Feature,
169        }
170    }
171
172    pub fn __str__(&self) -> String {
173        ProfileFuncs::__str__(self)
174    }
175}
176
177impl Features {
178    pub fn iter(&self) -> std::slice::Iter<'_, Feature> {
179        self.features.iter()
180    }
181}
182
183#[pyclass]
184#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)]
185pub struct FeatureMap {
186    #[pyo3(get)]
187    pub features: HashMap<String, HashMap<String, usize>>,
188}
189
190#[pymethods]
191impl FeatureMap {
192    pub fn __str__(&self) -> String {
193        // serialize the struct to a string
194        ProfileFuncs::__str__(self)
195    }
196}
197
198#[pyclass]
199#[derive(Clone, Serialize, Debug)]
200pub struct Metric {
201    pub name: String,
202    pub value: f64,
203}
204
205#[pymethods]
206impl Metric {
207    #[new]
208    pub fn new(name: String, value: f64) -> Self {
209        Metric { name, value }
210    }
211    pub fn __str__(&self) -> String {
212        ProfileFuncs::__str__(self)
213    }
214}
215
216#[pyclass]
217#[derive(Clone, Serialize, Debug)]
218pub struct Metrics {
219    #[pyo3(get)]
220    pub metrics: Vec<Metric>,
221
222    #[pyo3(get)]
223    pub entity_type: EntityType,
224}
225
226#[pymethods]
227impl Metrics {
228    #[new]
229    pub fn new(metrics: Vec<Metric>) -> Self {
230        Metrics {
231            metrics,
232            entity_type: EntityType::Metric,
233        }
234    }
235    pub fn __str__(&self) -> String {
236        ProfileFuncs::__str__(self)
237    }
238}
239
240impl Metrics {
241    pub fn iter(&self) -> std::slice::Iter<'_, Metric> {
242        self.metrics.iter()
243    }
244}
245
246#[derive(Debug)]
247pub enum QueueItem {
248    Features(Features),
249    Metrics(Metrics),
250}
251
252impl QueueItem {
253    /// Helper for extracting an Entity from a Python object
254    pub fn from_py_entity(entity: &Bound<'_, PyAny>) -> Result<Self, TypeError> {
255        let entity_type = entity.getattr("entity_type")?.extract::<EntityType>()?;
256
257        match entity_type {
258            EntityType::Feature => {
259                let features = entity.extract::<Features>()?;
260                Ok(QueueItem::Features(features))
261            }
262            EntityType::Metric => {
263                let metrics = entity.extract::<Metrics>()?;
264                Ok(QueueItem::Metrics(metrics))
265            }
266        }
267    }
268}
269
270pub trait QueueExt: Send + Sync {
271    fn metrics(&self) -> &Vec<Metric>;
272    fn features(&self) -> &Vec<Feature>;
273}
274
275impl QueueExt for Features {
276    fn metrics(&self) -> &Vec<Metric> {
277        // this is not a real implementation, just a placeholder
278        // to satisfy the trait bound
279        static EMPTY: Vec<Metric> = Vec::new();
280        &EMPTY
281    }
282
283    fn features(&self) -> &Vec<Feature> {
284        &self.features
285    }
286}
287
288impl QueueExt for Metrics {
289    fn metrics(&self) -> &Vec<Metric> {
290        &self.metrics
291    }
292
293    fn features(&self) -> &Vec<Feature> {
294        // this is not a real implementation, just a placeholder
295        // to satisfy the trait bound
296        static EMPTY: Vec<Feature> = Vec::new();
297        &EMPTY
298    }
299}
scouter_types/queue/types.rs

scouter_types/queue/
types.rs
