recursive.py

recursive.py

def change_type(root, old_type, new_type):
    """
    Helper function to recursively walk through the datastructure and change old_type to new_type
    """
    def _match(obj, **kwargs):
        return isinstance(obj, old_type)

    def _replace(obj, **kwargs):
        return new_type(obj)

    return modify(root, _match, _replace)


def modify(root, match_func, replace_func):
    """
    Apply recursively a function func to all leaves of a complex structure started with `root`

    :param match_func: function which returns True if object has to be modified
    :param replace_func: function which is called for every object, marked as "to replacement" by
                         match function.

    Both functions must have following signature::

        def my_func(obj, key, level):
            ...

    Here `obj` is the object to inspect / modify, `key` is the key name, if object is a part of the dict,
    and `level` is the depth level.
    """

    def _wrap(obj, key, level):

        if match_func(obj, key=key, level=level):
            obj = replace_func(obj, key=key, level=level)

        if isinstance(obj, dict):
            return obj.__class__({k: _wrap(v, k, level+1) for k, v in obj.items()})

        if isinstance(obj, (list, tuple, set)):
            return obj.__class__([_wrap(item, None, level+1) for item in obj])

        return obj

    return _wrap(root, None, 0)


def walk(root, match_func):
    """
    Walk through the structure and find all object of a certain type

    For example, return integer values from an object::

        >>> list(walk([1, 2, [3, 4, 5], {'foo': 6}], lambda obj, **kw: isinstance(obj, int)))
        [1, 2, 3, 4, 5, 6]
    """
    def _wrap(obj, key, level):

        if match_func(obj, key=key, level=level):
            yield obj

        if isinstance(obj, dict):
            for k, v in obj.items():
                for i in _wrap(v, k, level+1):
                    yield i

        if isinstance(obj, (list, tuple, set)):
            for item in obj:
                for i in _wrap(item, None, level+1):
                    yield i

    return _wrap(root, None, 0)