1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
use std::{convert::TryFrom, path::Path, sync::Arc};

use pyo3::{
    exceptions::{PyRuntimeError, PyTypeError},
    prelude::*,
    types::{PyByteArray, PyTuple},
};
use pythonize::pythonize;

use crate::{
    effects::{RuntimeMethodError, RuntimeMethods},
    image::RawImage,
    models::Color,
};

mod context;
use context::Context;

pub type Error = pyo3::PyErr;

#[derive(Default, Debug, Clone, Copy)]
pub struct PythonProvider;

impl PythonProvider {
    pub fn new() -> Self {
        Self
    }
}

impl super::Provider for PythonProvider {
    fn supports(&self, script_path: &str) -> bool {
        script_path.ends_with(".py")
    }

    fn run(
        &self,
        full_script_path: &Path,
        args: serde_json::Value,
        methods: Arc<dyn RuntimeMethods>,
    ) -> Result<(), super::ProviderError> {
        Ok(do_run(methods, args, |py| {
            // Run script
            py.run(
                std::fs::read_to_string(full_script_path)?.as_str(),
                None,
                None,
            )?;

            Ok(())
        })?)
    }
}

impl From<RuntimeMethodError> for PyErr {
    fn from(value: RuntimeMethodError) -> PyErr {
        match value {
            RuntimeMethodError::InvalidArguments { .. } => PyTypeError::new_err(value.to_string()),
            _ => PyRuntimeError::new_err(value.to_string()),
        }
    }
}

/// Check if the effect should abort execution
#[pyfunction]
fn abort() -> bool {
    Context::with_current(|m| async move { m.abort().await })
}

/// Set a new color for the leds
#[pyfunction(signature = (*args))]
#[pyo3(name = "setColor")]
fn set_color(args: &PyTuple) -> Result<(), PyErr> {
    Context::with_current(|m| {
        async move {
            if let Result::<(u8, u8, u8), _>::Ok((r, g, b)) = args.extract() {
                m.set_color(Color::new(r, g, b)).await?;
            } else if let Result::<(&PyByteArray,), _>::Ok((bytearray,)) = args.extract() {
                if bytearray.len() == 3 * m.get_led_count() {
                    // Safety: we are not modifying bytearray while accessing it
                    unsafe {
                        m.set_led_colors(
                            bytearray
                                .as_bytes()
                                .chunks_exact(3)
                                .map(|rgb| Color::new(rgb[0], rgb[1], rgb[2]))
                                .collect(),
                        )
                        .await?;
                    }
                } else {
                    return Err(RuntimeMethodError::InvalidByteArray.into());
                }
            } else {
                return Err(RuntimeMethodError::InvalidArguments { name: "setColor" }.into());
            }

            Ok(())
        }
    })
}

/// Set a new image to process and determine new led colors
#[pyfunction]
#[pyo3(name = "setImage")]
fn set_image(width: u16, height: u16, data: &PyByteArray) -> Result<(), PyErr> {
    Context::with_current(|m| {
        async move {
            // unwrap: we did all the necessary checks already
            m.set_image(
                RawImage::try_from((data.to_vec(), width as u32, height as u32))
                    .map_err(RuntimeMethodError::InvalidImageData)?,
            )
            .await?;

            Ok(())
        }
    })
}

#[pymodule]
fn hyperion(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(abort, m)?)?;
    m.add_function(wrap_pyfunction!(set_color, m)?)?;
    m.add_function(wrap_pyfunction!(set_image, m)?)?;

    m.add(
        "ledCount",
        Context::with_current(|m| async move { m.get_led_count() }),
    )?;

    Ok(())
}

extern "C" fn hyperion_init() -> *mut pyo3::ffi::PyObject {
    unsafe { hyperion::init() }
}

fn do_run<T>(
    methods: Arc<dyn RuntimeMethods>,
    args: serde_json::Value,
    f: impl FnOnce(Python) -> Result<T, PyErr>,
) -> Result<T, PyErr> {
    Context::with(methods, |ctx| {
        // Run the given code
        Python::with_gil(|py| {
            ctx.run(py, || {
                // Register arguments
                let hyperion_mod = py.import("hyperion")?;
                hyperion_mod.add("args", pythonize(py, &args)?)?;

                f(py)
            })
        })
    })
}

#[cfg(test)]
mod tests;