scouter_types/queue/
types.rs

1use crate::error::TypeError;
2use crate::ProfileFuncs;
3use pyo3::prelude::*;
4use pyo3::types::{PyDict, PyFloat, PyInt, PyList, PyString};
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    #[new]
92    /// Parses a value to it's corresponding feature type.
93    /// PyFLoat -> FloatFeature
94    /// PyInt -> IntFeature
95    /// PyString -> StringFeature
96    /// # Arguments
97    /// * `name` - The name of the feature.
98    /// * `feature` - The value of the feature, which can be a PyFloat
99    /// # Returns
100    /// * `Feature` - The corresponding feature type.
101    /// # Errors
102    /// * `TypeError` - If the feature type is not supported.
103    pub fn new(name: &str, feature: Bound<'_, PyAny>) -> Result<Self, TypeError> {
104        // check python type
105        if feature.is_instance_of::<PyFloat>() {
106            let value: f64 = feature.extract().unwrap();
107            Ok(Feature::Float(FloatFeature {
108                name: name.into(),
109                value,
110            }))
111        } else if feature.is_instance_of::<PyInt>() {
112            let value: i64 = feature.extract().unwrap();
113            Ok(Feature::Int(IntFeature {
114                name: name.into(),
115                value,
116            }))
117        } else if feature.is_instance_of::<PyString>() {
118            let value: String = feature.extract().unwrap();
119            Ok(Feature::String(StringFeature {
120                name: name.into(),
121                value,
122            }))
123        } else {
124            Err(TypeError::UnsupportedFeatureTypeError(
125                feature.get_type().name()?.to_string(),
126            ))
127        }
128    }
129
130    #[staticmethod]
131    pub fn int(name: String, value: i64) -> Self {
132        Feature::Int(IntFeature { name, value })
133    }
134
135    #[staticmethod]
136    pub fn float(name: String, value: f64) -> Self {
137        Feature::Float(FloatFeature { name, value })
138    }
139
140    #[staticmethod]
141    pub fn string(name: String, value: String) -> Self {
142        Feature::String(StringFeature { name, value })
143    }
144
145    #[staticmethod]
146    pub fn categorical(name: String, value: String) -> Self {
147        Feature::String(StringFeature { name, value })
148    }
149
150    pub fn __str__(&self) -> String {
151        ProfileFuncs::__str__(self)
152    }
153}
154
155impl Feature {
156    pub fn to_float(&self, feature_map: &FeatureMap) -> Result<f64, TypeError> {
157        match self {
158            Feature::Int(feature) => Ok(feature.to_float()),
159            Feature::Float(feature) => Ok(feature.value),
160            Feature::String(feature) => feature.to_float(feature_map),
161        }
162    }
163
164    pub fn name(&self) -> &str {
165        match self {
166            Feature::Int(feature) => &feature.name,
167            Feature::Float(feature) => &feature.name,
168            Feature::String(feature) => &feature.name,
169        }
170    }
171
172    pub fn to_usize(&self, feature_map: &FeatureMap) -> Result<usize, TypeError> {
173        match self {
174            Feature::Int(f) => Ok(f.value as usize),
175            Feature::Float(f) => Ok(f.value as usize),
176            Feature::String(f) => Ok(f.to_float(feature_map)? as usize),
177        }
178    }
179}
180
181impl Display for Feature {
182    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
183        match self {
184            Feature::Int(feature) => write!(f, "{}", feature.value),
185            Feature::Float(feature) => write!(f, "{}", feature.value),
186            Feature::String(feature) => write!(f, "{}", feature.value),
187        }
188    }
189}
190
191#[pyclass]
192#[derive(Clone, Debug, Serialize)]
193pub struct Features {
194    #[pyo3(get)]
195    pub features: Vec<Feature>,
196
197    #[pyo3(get)]
198    pub entity_type: EntityType,
199}
200
201#[pymethods]
202impl Features {
203    #[new]
204    /// Creates a new Features instance.
205    /// A user may supply either a list of features or a single feature.
206    /// Extract features into a Vec<Feature>
207    /// Extraction follows the following rules:
208    /// 1. Check if Pylist, if so, extract to Vec<Feature>
209    /// 2. Check if PyDict, if so, iterate over each key-value pair and create a Feature
210    /// 3. If neither, return an error
211    /// # Arguments
212    /// * `features` - A Python object that can be a list of Feature instances or
213    ///               a dictionary of key-value pairs where keys are feature names
214    /// # Returns
215    /// * `Features` - A new Features instance containing the extracted features.
216    pub fn new(features: Bound<'_, PyAny>) -> Result<Self, TypeError> {
217        let features = if features.is_instance_of::<PyList>() {
218            features
219                .downcast::<PyList>()
220                .unwrap()
221                .iter()
222                .map(|item| item.extract::<Feature>().unwrap())
223                .collect()
224        } else if features.is_instance_of::<PyDict>() {
225            features
226                .downcast::<PyDict>()
227                .unwrap()
228                .iter()
229                .map(|(key, value)| Feature::new(key.extract().unwrap(), value.clone()).unwrap())
230                .collect()
231        } else {
232            Err(TypeError::UnsupportedFeaturesTypeError(
233                features.get_type().name()?.to_string(),
234            ))?
235        };
236        Ok(Features {
237            features,
238            entity_type: EntityType::Feature,
239        })
240    }
241
242    pub fn __str__(&self) -> String {
243        ProfileFuncs::__str__(self)
244    }
245}
246
247impl Features {
248    pub fn iter(&self) -> std::slice::Iter<'_, Feature> {
249        self.features.iter()
250    }
251}
252
253#[pyclass]
254#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)]
255pub struct FeatureMap {
256    #[pyo3(get)]
257    pub features: HashMap<String, HashMap<String, usize>>,
258}
259
260#[pymethods]
261impl FeatureMap {
262    pub fn __str__(&self) -> String {
263        // serialize the struct to a string
264        ProfileFuncs::__str__(self)
265    }
266}
267
268#[pyclass]
269#[derive(Clone, Serialize, Debug)]
270pub struct Metric {
271    pub name: String,
272    pub value: f64,
273}
274
275#[pymethods]
276impl Metric {
277    #[new]
278    pub fn new(name: String, value: Bound<'_, PyAny>) -> Self {
279        let value = if value.is_instance_of::<PyFloat>() {
280            value.extract::<f64>().unwrap()
281        } else if value.is_instance_of::<PyInt>() {
282            value.extract::<i64>().unwrap() as f64
283        } else {
284            panic!(
285                "Unsupported metric type: {}",
286                value.get_type().name().unwrap()
287            );
288        };
289        Metric { name, value }
290    }
291
292    pub fn __str__(&self) -> String {
293        ProfileFuncs::__str__(self)
294    }
295}
296
297impl Metric {
298    pub fn new_rs(name: String, value: f64) -> Self {
299        Metric { name, value }
300    }
301}
302
303#[pyclass]
304#[derive(Clone, Serialize, Debug)]
305pub struct Metrics {
306    #[pyo3(get)]
307    pub metrics: Vec<Metric>,
308
309    #[pyo3(get)]
310    pub entity_type: EntityType,
311}
312
313#[pymethods]
314impl Metrics {
315    #[new]
316    pub fn new(metrics: Bound<'_, PyAny>) -> Result<Self, TypeError> {
317        let metrics = if metrics.is_instance_of::<PyList>() {
318            metrics
319                .downcast::<PyList>()
320                .unwrap()
321                .iter()
322                .map(|item| item.extract::<Metric>().unwrap())
323                .collect()
324        } else if metrics.is_instance_of::<PyDict>() {
325            metrics
326                .downcast::<PyDict>()
327                .unwrap()
328                .iter()
329                .map(|(key, value)| Metric::new(key.extract().unwrap(), value))
330                .collect()
331        } else {
332            Err(TypeError::UnsupportedMetricsTypeError(
333                metrics.get_type().name()?.to_string(),
334            ))?
335        };
336        Ok(Metrics {
337            metrics,
338            entity_type: EntityType::Metric,
339        })
340    }
341    pub fn __str__(&self) -> String {
342        ProfileFuncs::__str__(self)
343    }
344}
345
346impl Metrics {
347    pub fn iter(&self) -> std::slice::Iter<'_, Metric> {
348        self.metrics.iter()
349    }
350}
351
352#[derive(Debug)]
353pub enum QueueItem {
354    Features(Features),
355    Metrics(Metrics),
356}
357
358impl QueueItem {
359    /// Helper for extracting an Entity from a Python object
360    pub fn from_py_entity(entity: &Bound<'_, PyAny>) -> Result<Self, TypeError> {
361        let entity_type = entity.getattr("entity_type")?.extract::<EntityType>()?;
362
363        match entity_type {
364            EntityType::Feature => {
365                let features = entity.extract::<Features>()?;
366                Ok(QueueItem::Features(features))
367            }
368            EntityType::Metric => {
369                let metrics = entity.extract::<Metrics>()?;
370                Ok(QueueItem::Metrics(metrics))
371            }
372        }
373    }
374}
375
376pub trait QueueExt: Send + Sync {
377    fn metrics(&self) -> &Vec<Metric>;
378    fn features(&self) -> &Vec<Feature>;
379}
380
381impl QueueExt for Features {
382    fn metrics(&self) -> &Vec<Metric> {
383        // this is not a real implementation, just a placeholder
384        // to satisfy the trait bound
385        static EMPTY: Vec<Metric> = Vec::new();
386        &EMPTY
387    }
388
389    fn features(&self) -> &Vec<Feature> {
390        &self.features
391    }
392}
393
394impl QueueExt for Metrics {
395    fn metrics(&self) -> &Vec<Metric> {
396        &self.metrics
397    }
398
399    fn features(&self) -> &Vec<Feature> {
400        // this is not a real implementation, just a placeholder
401        // to satisfy the trait bound
402        static EMPTY: Vec<Feature> = Vec::new();
403        &EMPTY
404    }
405}
scouter_types/queue/types.rs

scouter_types/queue/
types.rs
