rustpython_vm/stdlib/
atexit.rs

1pub use atexit::_run_exitfuncs;
2pub(crate) use atexit::module_def;
3
4#[pymodule]
5mod atexit {
6    use crate::{
7        AsObject, PyObjectRef, PyResult, VirtualMachine, common::rc::PyRc, function::FuncArgs,
8    };
9
10    #[pyfunction]
11    fn register(func: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyObjectRef {
12        // Callbacks go in LIFO order (insert at front)
13        vm.state
14            .atexit_funcs
15            .lock()
16            .insert(0, PyRc::new((func.clone(), args)));
17        func
18    }
19
20    #[pyfunction]
21    fn _clear(vm: &VirtualMachine) {
22        vm.state.atexit_funcs.lock().clear();
23    }
24
25    #[pyfunction]
26    fn unregister(func: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
27        // Iterate backward (oldest to newest in LIFO list).
28        // Release the lock during comparison so __eq__ can call atexit functions.
29        let mut i = {
30            let funcs = vm.state.atexit_funcs.lock();
31            funcs.len() as isize - 1
32        };
33        while i >= 0 {
34            let entry = {
35                let funcs = vm.state.atexit_funcs.lock();
36                if i as usize >= funcs.len() {
37                    i = funcs.len() as isize;
38                    i -= 1;
39                    continue;
40                }
41                // Keep the entry alive for as long as it is being compared, so
42                // it cannot be dropped and have its address handed to a
43                // callback registered from within __eq__.
44                funcs[i as usize].clone()
45            };
46            // Lock released: __eq__ can safely call atexit functions
47            let eq = vm.bool_eq(&func, &entry.0)?;
48            if eq {
49                // The entry may have moved during __eq__. Search backward by identity.
50                let mut funcs = vm.state.atexit_funcs.lock();
51                let mut j = (funcs.len() as isize - 1).min(i);
52                while j >= 0 {
53                    if PyRc::ptr_eq(funcs.get(j as usize).unwrap(), &entry) {
54                        funcs.remove(j as usize);
55                        i = j;
56                        break;
57                    }
58                    j -= 1;
59                }
60            }
61            {
62                let funcs = vm.state.atexit_funcs.lock();
63                if i as usize >= funcs.len() {
64                    i = funcs.len() as isize;
65                }
66            }
67            i -= 1;
68        }
69        Ok(())
70    }
71
72    #[pyfunction]
73    pub fn _run_exitfuncs(vm: &VirtualMachine) {
74        let funcs: Vec<_> = core::mem::take(&mut *vm.state.atexit_funcs.lock());
75        // Callbacks stored in LIFO order, iterate forward
76        for entry in funcs {
77            let (func, args) = PyRc::try_unwrap(entry).unwrap_or_else(|e| (*e).clone());
78            if let Err(e) = func.call(args, vm) {
79                let exit = e.fast_isinstance(vm.ctx.exceptions.system_exit);
80                let msg = func
81                    .repr(vm)
82                    .ok()
83                    .map(|r| format!("Exception ignored in atexit callback {}", r.as_wtf8()));
84                vm.run_unraisable(e, msg, vm.ctx.none());
85                if exit {
86                    break;
87                }
88            }
89        }
90    }
91
92    #[pyfunction]
93    fn _ncallbacks(vm: &VirtualMachine) -> usize {
94        vm.state.atexit_funcs.lock().len()
95    }
96}
rustpython_vm/stdlib/atexit.rs

rustpython_vm/stdlib/
atexit.rs
