SZanlongo
7/28/2016 - 2:41 PM

ViviDict Creates and returns a nested dictionary that is either a regular dict or an LRU-style dict

ViviDict

Creates and returns a nested dictionary that is either a regular dict or an LRU-style dict

from collections import OrderedDict


def create_vividict(**kwds):
    """
    Creates and returns a nested dictionary that is either a regular dict or an LRU-style dict
    :param kwds: args
        :max_weighted_size: int
            The largest size the dict can have before it begins popping older cells
        :map_memory_weights: dict
            Dictionary holding the memory weights of each attribute
    :return: ViviDict
        Returns ViviDict object that is either a regular dict, or an LRU-style dict
    """
    ordering = kwds.pop("ordering")
    base = OrderedDict if ordering == "ordered" else dict
    max_weighted_size = kwds.pop("max_weighted_size", None)
    memory_weights = kwds.pop("memory_weights", {})

    class ViviDict(base):
        def __init__(self, *args):
            """
            ViviDict object
            :weighted_size: int
                Current weighted size of the dict
            :max_weighted_size: int
                 Max size the dictionary can have. If None, unlimited size
            :map_memory_weights: dict
                Dictionary with the memory weight for each attribute
            """
            self.weighted_size = 0
            base.__init__(self, *args)

        def __missing__(self, key):
            """
            If a value is missing, create a new ViviDict object in its place
            Based on: http://stackoverflow.com/q/635483
            :param key:
                Value to use as dict-key
            :return:
                ViviDict object as key-position
            """
            value = self[key] = type(self)()
            return value

        def set_value(self, keys, value):
            """
            Set the value at the given chain of keys
            Will replace existing value (if any)
            Will create the appropriate chain of keys if needed
            ex: [key1][key2][...] = value
            :param keys: list
                List of keys to travel down in dict
            :param value: value
                Value to set
            """
            # Check if LRU or not
            if self.__class__.__base__.__name__ == OrderedDict.__name__ and self.max_weighted_size is not None:
                # LRU-style, record the size of the attribute
                new_weight = 0
                if self.get_value(keys[:2]) is None:
                    # New attribute, add its weight (else, simply replacing a value, does not affect weight)
                    new_weight = self.memory_weights.get(keys[1], 0)
                    self.weighted_size += new_weight

                # Existing weight for the cell
                old_weight = max(0, self.get_value([keys[0], "weight"]))
                new_weight += old_weight
                # Set weight for the cell
                self.set_value_recursive(self, [keys[0], "weight"], new_weight)

                # Check if any older cells need to be removed to make space
                self.prune_dict()

            # Set value
            self.set_value_recursive(self, keys, value)

        def set_value_recursive(self, branch, keys, value):
            """
            Recursively travel down the list of keys in order to set the appropriate value
            :param branch: dict
                Dictionary to modify
            :param keys: list
                List of keys to travel down and modify or create as needed
            :param value: value
                Value to set at end of key-chain
            """
            if len(keys) > 1:
                if keys[0] not in branch:
                    # noinspection PyStatementEffect
                    branch[keys[0]]  # Create an empty ViviDict for next step
                # Recurse on remaining keys
                self.set_value_recursive(branch[keys[0]], keys[1:], value)
            else:
                # Final key, set its value
                branch[keys[0]] = value

        def get_value(self, keys):
            """
            Helper function to get the value of an attribute by travelling down the key-chain
            Helps in not creating new empty dicts when accessing a non-existent value
            :param keys: list
                List of keys to travel down
            :return: value if exists, else None
            """
            return self.get_value_recursive(self, keys)

        def get_value_recursive(self, branch, keys):
            """
            Get the value of an attribute by travelling down the key-chain
            Helps in not creating new empty dicts when accessing a non-existent value
            :return: value if exists, else None
            :param branch: dict
                Dictionary to check
            :param keys: list
                List of keys to travel down
            :return: value if exists, else None
            """
            if len(keys) > 1:
                if keys[0] in branch:
                    # Recurse on remaining keys
                    return self.get_value_recursive(branch[keys[0]], keys[1:])
                else:
                    # Key does not exist, return None
                    return None
            else:
                if keys[0] in branch:
                    # Final key, value exists, return
                    return branch[keys[0]]
                else:
                    # Key-value does not exist, return None
                    return None

        def delete(self, keys):
            """
            Helper function to delete a value if it exists
            Final key is the key whose value will be deleted
            Recursively checks if the parent-keys are now empty, and will prune those as needed
            :param keys: list
                List of keys to travel down
            """
            return self.delete_recursive(self, keys)

        def delete_recursive(self, branch, keys):
            """
            Delete a value if it exists
            Final key is the key whose value will be deleted
            Recursively checks if the parent-keys are now empty, and will prune those as needed
            Based on http://stackoverflow.com/q/21390243
            :param branch: dict
                Dictionary to modify
            :param keys:
                List of keys to travel down
            """
            if len(keys) > 1:  # not at the leaf
                empty = self.delete_recursive(branch[keys[0]], keys[1:])  # recursion
                if empty:
                    del branch[keys[0]]  # delete branch
            else:  # at the leaf
                if keys[0] in branch:
                    del branch[keys[0]]  # delete the leaf

            return len(branch) == 0  # could return len

        def prune_dict(self):
            """
            Prune dict to maintain total memory weight below max memory weight threshold
            """
            while self.weighted_size > self.max_weighted_size:
                # Keep deleting the oldest cell until weighted size is less than max weighted size threshold
                self.weighted_size -= self.popitem(last=False)[1]["weight"]

    vd = ViviDict()
    vd.max_weighted_size = max_weighted_size
    vd.memory_weights = memory_weights

    return vd