pite_script/
descriptor.rs1use 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}