@@ -3338,6 +3338,54 @@ pyeval_getlocals(PyObject *module, PyObject *Py_UNUSED(args))
33383338    return  Py_XNewRef (PyEval_GetLocals ());
33393339}
33403340
3341+ struct  atexit_data  {
3342+     int  called ;
3343+     PyThreadState  * tstate ;
3344+     PyInterpreterState  * interp ;
3345+ };
3346+ 
3347+ static  void 
3348+ atexit_callback (void  * data )
3349+ {
3350+     struct  atexit_data  * at_data  =  (struct  atexit_data  * )data ;
3351+     // Ensure that the callback is from the same interpreter 
3352+     assert (PyThreadState_Get () ==  at_data -> tstate );
3353+     assert (PyInterpreterState_Get () ==  at_data -> interp );
3354+     ++ at_data -> called ;
3355+ }
3356+ 
3357+ static  PyObject  * 
3358+ test_atexit (PyObject  * self , PyObject  * Py_UNUSED (args ))
3359+ {
3360+     PyThreadState  * oldts  =  PyThreadState_Swap (NULL );
3361+     PyThreadState  * tstate  =  Py_NewInterpreter ();
3362+ 
3363+     struct  atexit_data  data  =  {0 };
3364+     data .tstate  =  PyThreadState_Get ();
3365+     data .interp  =  PyInterpreterState_Get ();
3366+ 
3367+     int  amount  =  10 ;
3368+     for  (int  i  =  0 ; i  <  amount ; ++ i )
3369+     {
3370+         int  res  =  PyUnstable_AtExit (tstate -> interp , atexit_callback , (void  * )& data );
3371+         if  (res  <  0 ) {
3372+             Py_EndInterpreter (tstate );
3373+             PyThreadState_Swap (oldts );
3374+             PyErr_SetString (PyExc_RuntimeError , "atexit callback failed" );
3375+             return  NULL ;
3376+         }
3377+     }
3378+ 
3379+     Py_EndInterpreter (tstate );
3380+     PyThreadState_Swap (oldts );
3381+ 
3382+     if  (data .called  !=  amount ) {
3383+         PyErr_SetString (PyExc_RuntimeError , "atexit callback not called" );
3384+         return  NULL ;
3385+     }
3386+     Py_RETURN_NONE ;
3387+ }
3388+ 
33413389static  PyMethodDef  TestMethods [] =  {
33423390    {"set_errno" ,               set_errno ,                       METH_VARARGS },
33433391    {"test_config" ,             test_config ,                     METH_NOARGS },
@@ -3483,6 +3531,7 @@ static PyMethodDef TestMethods[] = {
34833531    {"function_set_warning" , function_set_warning , METH_NOARGS },
34843532    {"test_critical_sections" , test_critical_sections , METH_NOARGS },
34853533    {"pyeval_getlocals" , pyeval_getlocals , METH_NOARGS },
3534+     {"test_atexit" , test_atexit , METH_NOARGS },
34863535    {NULL , NULL } /* sentinel */ 
34873536};
34883537
0 commit comments