1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

pub use atexit::_run_exitfuncs;
pub(crate) use atexit::make_module;

#[pymodule]
mod atexit {
    use crate::{function::FuncArgs, AsObject, PyObjectRef, PyResult, VirtualMachine};

    #[pyfunction]
    fn register(func: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyObjectRef {
        vm.state.atexit_funcs.lock().push((func.clone(), args));
        func
    }

    #[pyfunction]
    fn _clear(vm: &VirtualMachine) {
        vm.state.atexit_funcs.lock().clear();
    }

    #[pyfunction]
    fn unregister(func: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
        let mut funcs = vm.state.atexit_funcs.lock();

        let mut i = 0;
        while i < funcs.len() {
            if vm.bool_eq(&funcs[i].0, &func)? {
                funcs.remove(i);
            } else {
                i += 1;
            }
        }

        Ok(())
    }

    #[pyfunction]
    pub fn _run_exitfuncs(vm: &VirtualMachine) {
        let funcs: Vec<_> = std::mem::take(&mut *vm.state.atexit_funcs.lock());
        for (func, args) in funcs.into_iter().rev() {
            if let Err(e) = func.call(args, vm) {
                let exit = e.fast_isinstance(vm.ctx.exceptions.system_exit);
                vm.run_unraisable(e, Some("Error in atexit._run_exitfuncs".to_owned()), func);
                if exit {
                    break;
                }
            }
        }
    }

    #[pyfunction]
    fn _ncallbacks(vm: &VirtualMachine) -> usize {
        vm.state.atexit_funcs.lock().len()
    }
}