walk_data - os.walk but for json-like data

walkdata.py

from unittest import TestCase


def walk_data(root):
    """Walk a nested data structure like os.walk"""
    data = root[-1]
    root = root[:-1]
    if hasattr(data, 'items'):
        for key, value in data.items():
            yield root + (key,), value
            yield from walk_data(root + (key, value))
    elif hasattr(data, '__delitem__'):
        for index, value in enumerate(data):
            yield root + (index,), value
            yield from walk_data(root + (index, value))


class WalkDataTestCase(TestCase):

    def test_object(self):
        data = {'foo': 'bar'}
        root = (data,)
        output = list(walk_data(root))
        self.assertEqual(output, [
            (('foo',), 'bar'),
        ])

    def test_list(self):
        data = ['foo', 'bar']
        root = (data,)
        output = list(walk_data(root))
        self.assertEqual(output, [
            ((0,), 'foo'),
            ((1,), 'bar'),
        ])

    def test_objects_recurse(self):
        data = {
            'foo': {
                'bar': 'blerg',
            },
        }
        root = (data,)
        output = list(walk_data(root))
        self.assertEqual(output, [
            (('foo',), {'bar': 'blerg'}),
            (('foo', 'bar'), 'blerg'),
        ])

    def test_list_recurse(self):
        data = {
            'foo': [
                {
                    'bar': 'blerg',
                },
            ],
        }
        root = (data,)
        output = list(walk_data(root))
        self.assertEqual(output, [
            (('foo',), [{'bar': 'blerg'}]),
            (('foo', 0), {'bar': 'blerg'}),
            (('foo', 0, 'bar'), 'blerg'),
        ])