rustpython_vm/stdlib/
atexit.rs

1pub use atexit::_run_exitfuncs;
2pub(crate) use atexit::make_module;
3
4#[pymodule]
5mod atexit {
6    use crate::{function::FuncArgs, AsObject, PyObjectRef, PyResult, VirtualMachine};
7
8    #[pyfunction]
9    fn register(func: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyObjectRef {
10        vm.state.atexit_funcs.lock().push((func.clone(), args));
11        func
12    }
13
14    #[pyfunction]
15    fn _clear(vm: &VirtualMachine) {
16        vm.state.atexit_funcs.lock().clear();
17    }
18
19    #[pyfunction]
20    fn unregister(func: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
21        let mut funcs = vm.state.atexit_funcs.lock();
22
23        let mut i = 0;
24        while i < funcs.len() {
25            if vm.bool_eq(&funcs[i].0, &func)? {
26                funcs.remove(i);
27            } else {
28                i += 1;
29            }
30        }
31
32        Ok(())
33    }
34
35    #[pyfunction]
36    pub fn _run_exitfuncs(vm: &VirtualMachine) {
37        let funcs: Vec<_> = std::mem::take(&mut *vm.state.atexit_funcs.lock());
38        for (func, args) in funcs.into_iter().rev() {
39            if let Err(e) = func.call(args, vm) {
40                let exit = e.fast_isinstance(vm.ctx.exceptions.system_exit);
41                vm.run_unraisable(e, Some("Error in atexit._run_exitfuncs".to_owned()), func);
42                if exit {
43                    break;
44                }
45            }
46        }
47    }
48
49    #[pyfunction]
50    fn _ncallbacks(vm: &VirtualMachine) -> usize {
51        vm.state.atexit_funcs.lock().len()
52    }
53}