How do I make an exit handler that runs after all weakref finalizers run

I have a module meant to be consumed as a library by other programs. I am using weakref.finalize to register cleanup functions for these objects, and atexit.register for some top level cleanup. I need all of the object finalizers to run before the top level cleanup.

The structure looks something like this:

#module.py
from weakref import finalize
import atexit

class MyObject:
    @staticmethod
    def _cleanup():
        print('Cleaning up module object')
    
    def __init__(self):
        finalize(self, MyObject._cleanup)

def _cleanup():
    print('Cleaning up module')

atexit.register(_cleanup)

Looking at the source code of weakref, weakref registers an atexit handler as soon as one finalizer is created. This means that, in the case of the example module, the top-level cleanup will run after all the finalizers only if nothing else has created a finalizer by the time module has been imported.

#user_code.py
from weakref import finalize

class UserObject:
    @staticmethod
    def _cleanup():
        print('Cleaning up user object')

    def __init__(self):
        finalize(self, UserObject._cleanup)

do_user_object = False
if do_user_object:
    u = UserObject()
   
import module
m = module.MyObject()

if do_user_object = True the output is:

Cleaning up module
Cleaning up module object
Cleaning up user object

if do_user_object = False the output is:

Cleaning up module object
Cleaning up module

So the cleanup order is dependent on user actions. How can I make the top-level cleanup always run after all of that module's objects are finalized?

Solution

weakref.finalize also uses atexit.register(). When the first instance of the finalize class is initialized, it checks if the class has been registered with atexit, and adds an entry for the classmethod finalize._exitfunc() if it has not. That class cleanup method then processes all of the finalize objects in the reverse order of how they were created.

Cleanup actions registered with atexit are also run in the reverse order of how they were registered.

As a result, if any other code calls weakref.finalize before your module is loaded, your cleanup function set with atexit.register will always be execute after all of the finalize code.

The solution is to instead log your module's cleanup code with weakref.finalize as well. If your module creates any instances of it's own classes that will have cleanup functions, than the module cleanup code should be sent to finalize before any of those objects are created. This will guarantee the module cleanup runs after the finalize code for each of your module objects. Here's an example of a valid module level call to finalize that is linked to the cleanup function itself:

# module.py
from weakref import finalize


class MyObject:
    @staticmethod
    def _cleanup():
        print('Cleaning up module object')

    def __init__(self):
        finalize(self, MyObject._cleanup)


def _cleanup():
    print('Cleaning up module')

finalize(_cleanup, _cleanup)

Then if do_user_object = True, the example above will get the output:

Cleaning up module object
Cleaning up module
Cleaning up user object

For reference of anyone that needs it, other object types that are valid for weakref are listed in the documentation here: https://docs.python.org/3/library/weakref.html