hyperion/effects/providers/
python.rs

1use std::{convert::TryFrom, ffi::CString, path::Path, sync::Arc};
2
3use pyo3::{
4    exceptions::{PyRuntimeError, PyTypeError},
5    prelude::*,
6    types::{PyByteArray, PyTuple},
7};
8use pythonize::pythonize;
9
10use crate::{
11    effects::{RuntimeMethodError, RuntimeMethods},
12    image::RawImage,
13    models::Color,
14};
15
16mod context;
17use context::Context;
18
19pub type Error = pyo3::PyErr;
20
21#[derive(Default, Debug, Clone, Copy)]
22pub struct PythonProvider;
23
24impl PythonProvider {
25    pub fn new() -> Self {
26        Self
27    }
28}
29
30impl super::Provider for PythonProvider {
31    fn supports(&self, script_path: &str) -> bool {
32        script_path.ends_with(".py")
33    }
34
35    fn run(
36        &self,
37        full_script_path: &Path,
38        args: serde_json::Value,
39        methods: Arc<dyn RuntimeMethods>,
40    ) -> Result<(), super::ProviderError> {
41        Ok(do_run(methods, args, |py| {
42            // Run script
43            py.run(
44                CString::new(std::fs::read_to_string(full_script_path)?)?.as_c_str(),
45                None,
46                None,
47            )?;
48
49            Ok(())
50        })?)
51    }
52}
53
54impl From<RuntimeMethodError> for PyErr {
55    fn from(value: RuntimeMethodError) -> PyErr {
56        match value {
57            RuntimeMethodError::InvalidArguments { .. } => PyTypeError::new_err(value.to_string()),
58            _ => PyRuntimeError::new_err(value.to_string()),
59        }
60    }
61}
62
63/// Check if the effect should abort execution
64#[pyfunction]
65fn abort() -> bool {
66    Context::with_current(|m| async move { m.abort().await })
67}
68
69/// Set a new color for the leds
70#[pyfunction(signature = (*args))]
71#[pyo3(name = "setColor")]
72fn set_color(args: Bound<'_, PyTuple>) -> Result<(), PyErr> {
73    Context::with_current(|m| {
74        async move {
75            if let Result::<(u8, u8, u8), _>::Ok((r, g, b)) = args.extract() {
76                m.set_color(Color::new(r, g, b)).await?;
77            } else if let Result::<(Bound<'_, PyByteArray>,), _>::Ok((bytearray,)) = args.extract()
78            {
79                if bytearray.len() == 3 * m.get_led_count() {
80                    // Safety: we are not modifying bytearray while accessing it
81                    unsafe {
82                        m.set_led_colors(
83                            bytearray
84                                .as_bytes()
85                                .chunks_exact(3)
86                                .map(|rgb| Color::new(rgb[0], rgb[1], rgb[2]))
87                                .collect(),
88                        )
89                        .await?;
90                    }
91                } else {
92                    return Err(RuntimeMethodError::InvalidByteArray.into());
93                }
94            } else {
95                return Err(RuntimeMethodError::InvalidArguments { name: "setColor" }.into());
96            }
97
98            Ok(())
99        }
100    })
101}
102
103/// Set a new image to process and determine new led colors
104#[pyfunction]
105#[pyo3(name = "setImage")]
106fn set_image(width: u16, height: u16, data: Bound<'_, PyByteArray>) -> Result<(), PyErr> {
107    Context::with_current(|m| {
108        async move {
109            // unwrap: we did all the necessary checks already
110            m.set_image(
111                RawImage::try_from((data.to_vec(), width as u32, height as u32))
112                    .map_err(RuntimeMethodError::InvalidImageData)?,
113            )
114            .await?;
115
116            Ok(())
117        }
118    })
119}
120
121#[pymodule]
122fn hyperion(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
123    m.add_function(wrap_pyfunction!(abort, m)?)?;
124    m.add_function(wrap_pyfunction!(set_color, m)?)?;
125    m.add_function(wrap_pyfunction!(set_image, m)?)?;
126
127    m.add(
128        "ledCount",
129        Context::with_current(|m| async move { m.get_led_count() }),
130    )?;
131
132    Ok(())
133}
134
135fn do_run<T>(
136    methods: Arc<dyn RuntimeMethods>,
137    args: serde_json::Value,
138    f: impl FnOnce(Python) -> Result<T, PyErr>,
139) -> Result<T, PyErr> {
140    Context::with(methods, |ctx| {
141        // Run the given code
142        Python::attach(|py| {
143            ctx.run(py, || {
144                // Register arguments
145                let hyperion_mod = py.import("hyperion")?;
146                hyperion_mod.add("args", pythonize(py, &args)?)?;
147
148                f(py)
149            })
150        })
151    })
152}
153
154#[cfg(test)]
155mod tests;