nautilus_common/python/
wrappers.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Registry of the strong references which keep registered components' Python wrappers alive.
17//!
18//! Component inners hold only a weak reference to their Python wrapper, so something must own the
19//! wrapper for as long as the component stays registered. This registry is that owner, and it is
20//! thread-local for the same reason the component and actor registries are: those hold
21//! `Rc<UnsafeCell<..>>` and cannot leave their thread, so a wrapper registry with wider visibility
22//! would desync from the registrations it shadows.
23
24use std::cell::RefCell;
25
26use ahash::AHashMap;
27use nautilus_model::identifiers::ComponentId;
28use pyo3::prelude::*;
29
30thread_local! {
31    static PYTHON_WRAPPERS: RefCell<AHashMap<ComponentId, Py<PyAny>>> =
32        RefCell::new(AHashMap::new());
33}
34
35/// Retains the strong reference which keeps the Python wrapper for `component_id` alive.
36pub fn retain_python_wrapper(component_id: ComponentId, wrapper: Py<PyAny>) {
37    let displaced =
38        PYTHON_WRAPPERS.with_borrow_mut(|wrappers| wrappers.insert(component_id, wrapper));
39
40    if displaced.is_some() {
41        log::warn!("Replaced the retained Python wrapper for {component_id}");
42    }
43
44    // Dropping a wrapper can run Python finalization which re-enters Rust, so the value leaves the
45    // registry before the borrow ends
46    drop(displaced);
47}
48
49/// Releases the strong reference retained for `component_id`.
50pub fn release_python_wrapper(component_id: ComponentId) {
51    let released = PYTHON_WRAPPERS.with_borrow_mut(|wrappers| wrappers.remove(&component_id));
52
53    drop(released);
54}
55
56/// Returns the Python wrapper retained for `component_id`, or `None` when nothing is retained.
57#[must_use]
58pub fn get_python_wrapper(component_id: ComponentId) -> Option<Py<PyAny>> {
59    Python::attach(|py| {
60        PYTHON_WRAPPERS.with_borrow(|wrappers| {
61            wrappers
62                .get(&component_id)
63                .map(|wrapper| wrapper.clone_ref(py))
64        })
65    })
66}
67
68#[cfg(test)]
69mod tests {
70    use pyo3::{ffi::c_str, types::PyModule, wrap_pyfunction};
71    use rstest::rstest;
72
73    use super::*;
74
75    /// Lets a finalizing Python wrapper re-enter the registry.
76    #[pyfunction]
77    fn wrapper_is_retained(component_id: &str) -> bool {
78        get_python_wrapper(ComponentId::from(component_id)).is_some()
79    }
80
81    #[rstest]
82    fn test_wrapper_finalization_re_enters_an_unborrowed_registry() {
83        Python::initialize();
84
85        Python::attach(|py| {
86            let module = PyModule::new(py, "test_wrapper_finalization").unwrap();
87            module
88                .add_function(wrap_pyfunction!(wrapper_is_retained, &module).unwrap())
89                .unwrap();
90
91            let code = c_str!(
92                r#"
93OBSERVED = []
94
95
96class Finalizing:
97    def __del__(self):
98        OBSERVED.append(wrapper_is_retained("Finalizing-Component"))
99"#
100            );
101            py.run(code, Some(&module.dict()), None).unwrap();
102
103            let finalizing = module.getattr("Finalizing").unwrap();
104            let component_id = ComponentId::from("Finalizing-Component");
105
106            retain_python_wrapper(component_id, finalizing.call0().unwrap().unbind());
107
108            // The registry holds the only reference to each wrapper, so both the displaced wrapper
109            // and the released one run `__del__` while being dropped
110            retain_python_wrapper(component_id, finalizing.call0().unwrap().unbind());
111            release_python_wrapper(component_id);
112
113            let observed = module
114                .getattr("OBSERVED")
115                .unwrap()
116                .extract::<Vec<bool>>()
117                .unwrap();
118
119            assert_eq!(observed, vec![true, false]);
120        });
121    }
122}
nautilus_common/python/wrappers.rs

nautilus_common/python/
wrappers.rs
