little util to grab all values from collection matching a key

get_all_key_values.py

import unittest


def get_all_key_values(
    key,
    container
):
    result_set = set()
    if not isiterable(container):
        return result_set

    current_containers = [container]
    next_containers = []

    # avoid hanging
    MAX_RUNS = 100000
    runs = 0

    while True:
        for cur_container in current_containers:
            if isinstance(cur_container, list) and len(cur_container) > 0:
                current_containers.extend(cur_container)
                runs += 1
                continue
            for c_key in cur_container:
                container_val = cur_container[c_key]
                if isiterable(container_val):
                    next_containers.append(container_val)
                elif c_key == key:
                    result_set.add(container_val)

        if not next_containers:
            break


        current_containers = next_containers
        next_containers = []
        runs += 1

        if runs > MAX_RUNS:
            break

    return result_set


def isiterable(v):
    return isinstance(v, (dict, list))


class GetAllValuesForKeyTests(unittest.TestCase):
    def test_it_returns_empty_list_if_non_collection_container(self):
        values = get_all_key_values(
            key='id',
            container=1
        )
        self.assertEqual(set(), values)

    def test_shallow(self):
        values = get_all_key_values(
            key='id',
            container={
                'id': 2
            }
        )
        self.assertEqual({2}, values)

    def test_nested(self):
        values = get_all_key_values(
            key='id',
            container={
                'id': 1,
                'non-id1': -1,
                'nested-key': {
                    'id': 2,
                    'non-id2': -2
                }
            }
        )
        expected_members = [1, 2]
        for expected_member in expected_members:
            self.assertIn(expected_member, values)

    def test_with_list(self):
        values = get_all_key_values(
            key='id',
            container=[
                {'id': 1, 'not-id1': -1},
                {'id': 2, 'not-id1': -2}
            ]
        )
        expected_members = [1, 2]
        for expected_member in expected_members:
            self.assertIn(expected_member, values)

    def test_deep_nested_with_list(self):
        values = get_all_key_values(
            key='id',
            container={
                'id': 1,
                'nested-key1': {
                    'id': 2,
                    'items': [
                        {'id': 3, 'not-id1': -1},
                        {'id': 4, 'not-id1': -2}
                    ]
                }
            }
        )
        expected_members = [1, 2, 3, 4]
        for expected_member in expected_members:
            self.assertIn(expected_member, values)

    def test_it_ignores_when_key_is_compound_type(self):
        values = get_all_key_values(
            key='id',
            container={
                'id': {}
            }
        )
        for value in values:
            self.assertEqual({}, value)

    def test_it_returns_unique_values(self):
        values = get_all_key_values(
            key='id',
            container={
                'id': 1,
                'nested-key1': {
                    'id': 2,
                    'items': [
                        {'id': 3, 'not-id1': -1},
                        {'id': 3, 'not-id1': -2}
                    ]
                }
            }
        )
        expected = {1, 2, 3}
        self.assertEqual(expected, values)