Skip to main content

nyx_space/dynamics/sequence/
python.rs

1use pyo3::exceptions::PyException;
2use {
3    super::{SpacecraftSequence, Thruster},
4    crate::Spacecraft,
5    anise::almanac::Almanac,
6    pyo3::prelude::*,
7    std::{collections::HashMap, sync::Arc},
8};
9
10#[cfg(feature = "python")]
11#[pymethods]
12impl SpacecraftSequence {
13    #[new]
14    fn py_new() -> Self {
15        SpacecraftSequence::default()
16    }
17
18    #[classmethod]
19    #[pyo3(name = "from_dhall")]
20    fn py_from_dhall(_cls: &Bound<'_, pyo3::types::PyType>, dhall_str: &str) -> PyResult<Self> {
21        serde_dhall::from_str(dhall_str)
22            .parse()
23            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
24    }
25
26    #[classmethod]
27    #[pyo3(name = "from_yaml")]
28    fn py_from_yaml(_cls: &Bound<'_, pyo3::types::PyType>, yaml_str: &str) -> PyResult<Self> {
29        serde_yml::from_str(yaml_str)
30            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
31    }
32
33    #[pyo3(name = "setup")]
34    fn py_setup(&mut self, py: Python<'_>, almanac: Py<Almanac>) -> PyResult<()> {
35        let almanac_ref = almanac.borrow(py);
36        self.setup(Arc::new(almanac_ref.clone()))
37            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
38    }
39
40    #[pyo3(name = "propagate")]
41    fn py_propagate(
42        &self,
43        py: Python<'_>,
44        state: Spacecraft,
45        until_phase: Option<String>,
46        almanac: Py<Almanac>,
47    ) -> PyResult<Vec<(Option<String>, Vec<Spacecraft>)>> {
48        let almanac_ref = almanac.borrow(py);
49        let trajs = self
50            .propagate(state, until_phase, Arc::new(almanac_ref.clone()))
51            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
52
53        let result = trajs
54            .into_iter()
55            .map(|traj| (traj.name, traj.states))
56            .collect();
57        Ok(result)
58    }
59
60    #[getter]
61    fn get_thruster_sets(&self) -> HashMap<String, Thruster> {
62        self.thruster_sets.clone()
63    }
64
65    fn thruster_set_insert(&mut self, name: String, thruster: Thruster) {
66        self.thruster_sets.insert(name, thruster);
67    }
68    fn thruster_set_remove(&mut self, name: String) -> PyResult<()> {
69        if self.thruster_sets.remove(&name).is_none() {
70            Err(PyException::new_err(format!("{name} not in thruster set")))
71        } else {
72            Ok(())
73        }
74    }
75}