potato_util/
utils.rs

1use crate::error::UtilError;
2
3use colored_json::{Color, ColorMode, ColoredFormatter, PrettyFormatter, Styler};
4use pyo3::prelude::*;
5
6use pyo3::types::{
7    PyAny, PyBool, PyDict, PyDictMethods, PyFloat, PyInt, PyList, PyString, PyTuple,
8};
9use pyo3::IntoPyObjectExt;
10use serde::Serialize;
11use serde_json::json;
12use serde_json::Value;
13use std::path::Path;
14use uuid::Uuid;
15
16pub fn create_uuid7() -> String {
17    Uuid::now_v7().to_string()
18}
19
20pub struct PyHelperFuncs {}
21
22impl PyHelperFuncs {
23    pub fn __str__<T: Serialize>(object: T) -> String {
24        match ColoredFormatter::with_styler(
25            PrettyFormatter::default(),
26            Styler {
27                key: Color::Rgb(75, 57, 120).foreground(),
28                string_value: Color::Rgb(4, 205, 155).foreground(),
29                float_value: Color::Rgb(4, 205, 155).foreground(),
30                integer_value: Color::Rgb(4, 205, 155).foreground(),
31                bool_value: Color::Rgb(4, 205, 155).foreground(),
32                nil_value: Color::Rgb(4, 205, 155).foreground(),
33                ..Default::default()
34            },
35        )
36        .to_colored_json(&object, ColorMode::On)
37        {
38            Ok(json) => json,
39            Err(e) => format!("Failed to serialize to json: {e}"),
40        }
41        // serialize the struct to a string
42    }
43
44    pub fn __json__<T: Serialize>(object: T) -> String {
45        match serde_json::to_string_pretty(&object) {
46            Ok(json) => json,
47            Err(e) => format!("Failed to serialize to json: {e}"),
48        }
49    }
50
51    /// Save a struct to a JSON file
52    ///
53    /// # Arguments
54    ///
55    /// * `model` - A reference to a struct that implements the `Serialize` trait
56    /// * `path` - A reference to a `Path` object that holds the path to the file
57    ///
58    /// # Returns
59    ///
60    /// A `Result` containing `()` or a `UtilError`
61    ///
62    /// # Errors
63    ///
64    /// This function will return an error if:
65    /// - The struct cannot be serialized to a string
66    pub fn save_to_json<T>(model: T, path: &Path) -> Result<(), UtilError>
67    where
68        T: Serialize,
69    {
70        // serialize the struct to a string
71        let json =
72            serde_json::to_string_pretty(&model).map_err(|_| UtilError::SerializationError)?;
73
74        // ensure .json extension
75        let path = path.with_extension("json");
76
77        if !path.exists() {
78            // ensure path exists, create if not
79            let parent_path = path.parent().ok_or(UtilError::GetParentPathError)?;
80
81            std::fs::create_dir_all(parent_path).map_err(|_| UtilError::CreateDirectoryError)?;
82        }
83
84        std::fs::write(path, json).map_err(|_| UtilError::WriteError)?;
85
86        Ok(())
87    }
88}
89
90pub fn json_to_pyobject<'py>(
91    py: Python,
92    value: &Value,
93    dict: &Bound<'py, PyDict>,
94) -> Result<Bound<'py, PyDict>, UtilError> {
95    match value {
96        Value::Object(map) => {
97            for (k, v) in map {
98                let py_value = match v {
99                    Value::Null => py.None(),
100                    Value::Bool(b) => b.into_py_any(py)?,
101                    Value::Number(n) => {
102                        if let Some(i) = n.as_i64() {
103                            i.into_py_any(py)?
104                        } else if let Some(f) = n.as_f64() {
105                            f.into_py_any(py)?
106                        } else {
107                            return Err(UtilError::InvalidNumber);
108                        }
109                    }
110                    Value::String(s) => s.into_py_any(py)?,
111                    Value::Array(arr) => {
112                        let py_list = PyList::empty(py);
113                        for item in arr {
114                            let py_item = json_to_pyobject_value(py, item)?;
115                            py_list.append(py_item)?;
116                        }
117                        py_list.into_py_any(py)?
118                    }
119                    Value::Object(_) => {
120                        let nested_dict = PyDict::new(py);
121                        json_to_pyobject(py, v, &nested_dict)?;
122                        nested_dict.into_py_any(py)?
123                    }
124                };
125                dict.set_item(k, py_value)?;
126            }
127        }
128        _ => return Err(UtilError::RootMustBeObjectError),
129    }
130
131    Ok(dict.clone())
132}
133
134pub fn json_to_pyobject_value(py: Python, value: &Value) -> Result<PyObject, UtilError> {
135    Ok(match value {
136        Value::Null => py.None(),
137        Value::Bool(b) => b.into_py_any(py)?,
138        Value::Number(n) => {
139            if let Some(i) = n.as_i64() {
140                i.into_py_any(py)?
141            } else if let Some(f) = n.as_f64() {
142                f.into_py_any(py)?
143            } else {
144                return Err(UtilError::InvalidNumber);
145            }
146        }
147        Value::String(s) => s.into_py_any(py)?,
148        Value::Array(arr) => {
149            let py_list = PyList::empty(py);
150            for item in arr {
151                let py_item = json_to_pyobject_value(py, item)?;
152                py_list.append(py_item)?;
153            }
154            py_list.into_py_any(py)?
155        }
156        Value::Object(_) => {
157            let nested_dict = PyDict::new(py);
158            json_to_pyobject(py, value, &nested_dict)?;
159            nested_dict.into_py_any(py)?
160        }
161    })
162}
163
164pub fn pyobject_to_json(obj: &Bound<'_, PyAny>) -> Result<Value, UtilError> {
165    if obj.is_instance_of::<PyDict>() {
166        let dict = obj.downcast::<PyDict>()?;
167        let mut map = serde_json::Map::new();
168        for (key, value) in dict.iter() {
169            let key_str = key.extract::<String>()?;
170            let json_value = pyobject_to_json(&value)?;
171            map.insert(key_str, json_value);
172        }
173        Ok(Value::Object(map))
174    } else if obj.is_instance_of::<PyList>() {
175        let list = obj.downcast::<PyList>()?;
176        let mut vec = Vec::new();
177        for item in list.iter() {
178            vec.push(pyobject_to_json(&item)?);
179        }
180        Ok(Value::Array(vec))
181    } else if obj.is_instance_of::<PyTuple>() {
182        let tuple = obj.downcast::<PyTuple>()?;
183        let mut vec = Vec::new();
184        for item in tuple.iter() {
185            vec.push(pyobject_to_json(&item)?);
186        }
187        Ok(Value::Array(vec))
188    } else if obj.is_instance_of::<PyString>() {
189        let s = obj.extract::<String>()?;
190        Ok(Value::String(s))
191    } else if obj.is_instance_of::<PyFloat>() {
192        let f = obj.extract::<f64>()?;
193        Ok(json!(f))
194    } else if obj.is_instance_of::<PyBool>() {
195        let b = obj.extract::<bool>()?;
196        Ok(json!(b))
197    } else if obj.is_instance_of::<PyInt>() {
198        let i = obj.extract::<i64>()?;
199        Ok(json!(i))
200    } else if obj.is_none() {
201        Ok(Value::Null)
202    } else {
203        // display "cant show" for unsupported types
204        // call obj.str to get the string representation
205        // if error, default to "unsupported type"
206        let obj_str = match obj.str() {
207            Ok(s) => s
208                .extract::<String>()
209                .unwrap_or_else(|_| "unsupported type".to_string()),
210            Err(_) => "unsupported type".to_string(),
211        };
212
213        Ok(Value::String(obj_str))
214    }
215}
216
217pub fn version() -> String {
218    env!("CARGO_PKG_VERSION").to_string()
219}