rustpython_vm/stdlib/
atexit.rs

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

rustpython_vm/stdlib/
atexit.rs
