ordered_dict.py

from collections import OrderedDict
from contextlib import contextmanager


@contextmanager
def preserve_dict_order(d: OrderedDict):
    """Allow mutations on an OrderedDict, restoring its original order after
    """
    original_order = list(d)

    yield d

    old_keys = set(original_order)
    added_keys = [k for k in d if k not in old_keys]
    new_order = original_order + added_keys

    for key in new_order:
        if key in d:
            d[key] = d[key]