Skip to main content

pite_script/
descriptor.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2use pyo3::prelude::*;
3
4#[pyclass(name = "ChildDescriptor")]
5pub struct PyChildDescriptor {
6    name: String,
7    attr: Option<String>,
8}
9
10#[pymethods]
11impl PyChildDescriptor {
12    #[new]
13    fn new(name: String) -> Self {
14        Self { name, attr: None }
15    }
16
17    fn __set_name__(&mut self, _owner: &Bound<'_, PyAny>, name: String) {
18        self.attr = Some(name);
19    }
20
21    fn __get__(
22        &self,
23        py: Python<'_>,
24        obj: Option<&Bound<'_, PyAny>>,
25        _owner: Option<&Bound<'_, PyAny>>,
26    ) -> PyResult<Py<PyAny>> {
27        let Some(obj) = obj else {
28            return Ok(Bound::new(
29                py,
30                PyChildDescriptor {
31                    name: self.name.clone(),
32                    attr: self.attr.clone(),
33                },
34            )?
35            .into_any()
36            .unbind());
37        };
38
39        let cache = match obj.getattr("_descriptor_cache") {
40            Ok(c) => c,
41            Err(_) => {
42                let dict = pyo3::types::PyDict::new(py);
43                obj.setattr("_descriptor_cache", &dict)?;
44                dict.into_any()
45            }
46        };
47
48        let attr_key = self.attr.as_deref().unwrap_or(self.name.as_str());
49        if let Ok(cached) = cache.get_item(attr_key) {
50            return Ok(cached.into_any().unbind());
51        }
52
53        let resolved = obj.call_method1("get_node", (&self.name,))?;
54        cache.set_item(attr_key, &resolved)?;
55        Ok(resolved.into())
56    }
57}
58
59#[pyclass(name = "NodeDescriptor")]
60pub struct PyNodeDescriptor {
61    path: String,
62    attr: Option<String>,
63}
64
65#[pymethods]
66impl PyNodeDescriptor {
67    #[new]
68    fn new(path: String) -> Self {
69        Self { path, attr: None }
70    }
71
72    fn __set_name__(&mut self, _owner: &Bound<'_, PyAny>, name: String) {
73        self.attr = Some(name);
74    }
75
76    fn __get__(
77        &self,
78        py: Python<'_>,
79        obj: Option<&Bound<'_, PyAny>>,
80        _owner: Option<&Bound<'_, PyAny>>,
81    ) -> PyResult<Py<PyAny>> {
82        let Some(obj) = obj else {
83            return Ok(Bound::new(
84                py,
85                PyNodeDescriptor {
86                    path: self.path.clone(),
87                    attr: self.attr.clone(),
88                },
89            )?
90            .into_any()
91            .unbind());
92        };
93
94        let cache = match obj.getattr("_descriptor_cache") {
95            Ok(c) => c,
96            Err(_) => {
97                let dict = pyo3::types::PyDict::new(py);
98                obj.setattr("_descriptor_cache", &dict)?;
99                dict.into_any()
100            }
101        };
102
103        let attr_key = self.attr.as_deref().unwrap_or(self.path.as_str());
104        if let Ok(cached) = cache.get_item(attr_key) {
105            return Ok(cached.into_any().unbind());
106        }
107
108        let resolved = obj.call_method1("get_node", (&self.path,))?;
109        cache.set_item(attr_key, &resolved)?;
110        Ok(resolved.into())
111    }
112}
113
114#[pyfunction]
115pub fn child(name: String) -> PyChildDescriptor {
116    PyChildDescriptor::new(name)
117}
118
119#[pyfunction]
120pub fn node(path: String) -> PyNodeDescriptor {
121    PyNodeDescriptor::new(path)
122}