test_special_dict.py

import collections


class PermOriginDict(dict):
    """
    A special dict used for tracking permissions along with their origins.

    `__setitem__` has been overridden to expect a tuple(perm, origin)
    `__getitem__` will return only the perm
    `.perm_origin_stack` will return the stack of (perm, origin) set per key

    >>> perms = PermOriginDict()
    >>> perms['resource'] = 'read', 'default', 1
    >>> perms['resource']
    'read'
    >>> perms['resource'] = 'write', 'admin', 2
    >>> perms['resource']
    'write'
    >>> perms.perm_origin_stack
    {'resource': [('read', 'default', 1), ('write', 'admin', 2)]}
    """

    def __init__(self, *args, **kw):
        dict.__init__(self, *args, **kw)
        super(PermOriginDict, self).__setattr__('perm_origin_stack', collections.OrderedDict())

    def __setitem__(self, key, perm_origin_obj_id):
        print(f"setting: {perm_origin_obj_id}")
        if isinstance(perm_origin_obj_id, tuple):
            (perm, origin, obj_id) = perm_origin_obj_id
            self.perm_origin_stack.setdefault(key, []).append((perm, origin, obj_id))
        else:
            perm = perm_origin_obj_id
        dict.__setitem__(self, key, perm)

    # def __getstate__(self):
    #     dict_state = self.__dict__.copy()
    #     perm_origin_stack = self.perm_origin_stack.__dict__.copy()
    #     print(dict_state, perm_origin_stack)
    #     return dict_state, perm_origin_stack
    #
    # def __setstate__(self, state):
    #     print('setstate called')
    #     print(f'setstate os {state}')
    #     self.perm_origin_stack = collections.OrderedDict()


dd = PermOriginDict()

dd['key1'] = 'perm1', 'from_where1', 121
dd['key2'] = 'perm2', 'from_where2', 122
dd['key3'] = 'perm3', 'from_where3', 123
dd['key3'] = 'perm4', 'from_where4', 125


print(list(dd.keys()))
print(f"key:{dd['key1']}")
print(dd.perm_origin_stack)

print ('-- before pickle --')

import pickle

data = pickle.dumps(dd, protocol=pickle.HIGHEST_PROTOCOL)
print(data)

new_dd = pickle.loads(data)

print ('-- after pickle --')
print(list(dd.keys()))
print(f"key:{dd['key1']}")
print(dd.perm_origin_stack)