import inspect


def get_default_hook(obj, default_hooks):
    """Finds the default save/load hook to use with the given object.

    Follows the Method Resolution Order, i.e., if no hook is registered for
    the class of the object itself, also searches classes which the object
    inherits from.

    Arguments
    ---------
    obj : instance
        Instance of a class.
    default_hooks : dict
        Mapping from classes to (checkpointing hook) functions.

    Returns
    -------
    The correct method or None if no method is registered.

    Example
    -------
    >>> a = torch.nn.Module()
    >>> get_default_hook(a, DEFAULT_SAVE_HOOKS) == torch_save
    True
    """
    mro = inspect.getmro(type(obj))
    for cls in mro:
        if cls in default_hooks:
            return default_hooks[cls]
    # If we got here, no hook found
    return None