oxmpl-py 0.6.0

Python Bindings for the OxMPL Library
Documentation 
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

// Copyright (c) 2025 Junior Sundar
//
// SPDX-License-Identifier: BSD-3-Clause

use pyo3::{prelude::*, types::PyList};
use std::{any::Any, rc::Rc, sync::Arc};

use oxmpl::base::state::{
    CompoundState as OxmplCompoundState, RealVectorState as OxmplRealVectorState,
    SO2State as OxmplSO2State, SO3State as OxmplSO3State, State,
};

use crate::base::{PyRealVectorState, PySO2State, PySO3State};

/// A state that is composed of one or more other states.
///
/// This is useful for representing complex states, such as the state of a robot with multiple
/// joints, or a rigid body in space (which is composed of a translation and a rotation).
///
/// Args:
///     components (List[State]): A list of state objects (e.g., `RealVectorState`, `SO2State`).
#[pyclass(name = "CompoundState", unsendable)]
#[derive(Clone)]
pub struct PyCompoundState(pub Rc<OxmplCompoundState>);

#[pymethods]
impl PyCompoundState {
    #[new]
    fn new(components: Vec<PyObject>) -> PyResult<Self> {
        let mut rust_components: Vec<Box<dyn State>> = Vec::with_capacity(components.len());

        Python::with_gil(|py| {
            for comp_object in components {
                let comp_any = comp_object.bind(py);
                if let Ok(rv_state) = comp_any.extract::<PyRef<PyRealVectorState>>() {
                    rust_components.push(Box::new((*rv_state.0).clone()));
                } else if let Ok(so2_state) = comp_any.extract::<PyRef<PySO2State>>() {
                    rust_components.push(Box::new((*so2_state.0).clone()));
                } else if let Ok(so3_state) = comp_any.extract::<PyRef<PySO3State>>() {
                    rust_components.push(Box::new((*so3_state.0).clone()));
                } else {
                    return Err(pyo3::exceptions::PyValueError::new_err(format!(
                        "Object of type '{}' is not a valid state component.",
                        comp_any.get_type().name()?
                    )));
                }
            }
            Ok(())
        })?;
        let compound_state = OxmplCompoundState {
            components: rust_components,
        };
        Ok(Self(Rc::new(compound_state)))
    }

    /// list[State]: The list of component states.
    #[getter]
    fn get_components(&self, py: Python<'_>) -> PyResult<PyObject> {
        let list = PyList::empty(py);

        for component in &self.0.components {
            let component_any = &**component as &dyn Any;

            if let Some(rv_state) = component_any.downcast_ref::<OxmplRealVectorState>() {
                list.append(Py::new(py, PyRealVectorState(Arc::new(rv_state.clone())))?)?;
            } else if let Some(so2_state) = component_any.downcast_ref::<OxmplSO2State>() {
                list.append(Py::new(py, PySO2State(Arc::new(so2_state.clone())))?)?;
            } else if let Some(so3_state) = component_any.downcast_ref::<OxmplSO3State>() {
                list.append(Py::new(py, PySO3State(Arc::new(so3_state.clone())))?)?;
            } else {
                return Err(pyo3::exceptions::PyTypeError::new_err(
                    "Encountered an unknown state type when getting components.",
                ));
            }
        }

        Ok(list.into())
    }

    /// The number of component states.
    fn __len__(&self) -> usize {
        self.0.components.len()
    }

    fn __repr__(&self) -> String {
        format!(
            "<CompoundState with {} components>",
            self.0.components.len()
        )
    }
}