I need to write a function to count numbers of nests in an iterable. I reached my goal but wondering if there's any more neat/faster solution (maybe even a library or so?). Besides I'm afraid that my code is error prone.
Let's analyze two examples:
dct_test = {
"a": {1: 2, 2: 3, 3: 4},
"b": {"x": "y", "z": [1, 2, 3, {7, 8}]},
"c": {"l": (2, 3, 5), "j": {"a": {"k": 1, "l": [1, 3, (2, 7, {"x": 1})]}}},
"d": [1, (2, 3)]
}
tpl_test = (1, [3, 4, {5, 6, 7}], {"a": 5, "b": [9, 8]})
This is a visualization of what I'm trying to achieve:
Code:
from collections.abc import Iterable
from itertools import chain
# EXAMPLES
dct_test = {
"a": {1: 2, 2: 3, 3: 4},
"b": {"x": "y", "z": [1, 2, 3, {7, 8}]},
"c": {"l": (2, 3, 5), "j": {"a": {"k": 1, "l": [1, 3, (2, 7, {"x": 1})]}}},
"d": [1, (2, 3)]
}
tpl_test = (1, [3, 4, {5, 6, 7}], {"a": 5, "b": [9, 8]})
# FUNCTIONS
# checking if an argument is a 'real' iterable (list, tuple, set, dictionary)
def IsRealIterable(
iter: Iterable
):
return True if isinstance(iter, Iterable) and not isinstance(iter, str) else False
# converting an iterable to a list
def ConvertIterableToList(
iter: Iterable,
dct_keys=False, # False returns dict.values(), True returns dict.keys()
):
if IsRealIterable(iter):
if isinstance(iter, tuple) or isinstance(iter, set):
return list(iter)
elif isinstance(iter, dict):
return list(iter.keys()) if dct_keys else list(iter.values())
elif isinstance(iter, list):
return iter
else:
raise TypeError("Given argument is not a list, tuple, set or dictionary")
# returns a list of iterables extracted from an iterable
def ExtractFromIterable(
iter: Iterable
):
if IsRealIterable(iter):
return [i for i in ConvertIterableToList(iter) if IsRealIterable(i)]
else:
raise TypeError("Given argument is not a list, tuple, set or dictionary")
# returns number of levels in an iterable
def CountNestingLevels(
iter: Iterable,
cntr: int=0,
only_sublevels: bool=True # if False then returns number of all nesting levels - not only a number of nests
):
lst_to_inspect = ExtractFromIterable(iter)
lst_deep = [] # empty list for nested iterables
# analyzing structure
# if lst_to_inspect is empty then returns already counted number of levels
if not lst_to_inspect:
return cntr if only_sublevels else cntr + 1
else:
cntr += 1
for i in lst_to_inspect:
# temporary list for iterables from each element in lst_to_inspect
temp_list = ExtractFromIterable(i)
# elements (only non-empty) should be added to lst_deep because otherwise ExtractFromIterable() would overwrite existing values
lst_deep.append(temp_list) if temp_list else None
# flattening lst_deep. it contains lists returned by ExtractFromIterable so lst_deep has structure: [[[], [], ...]]
lst_deep = list(chain.from_iterable(lst_deep))
# using ExtractFromIterable in lst_deep may return the same structure as lst_to_inspect
# if no - use recursion to analyze remaining elements again
# if yes - finish analyzing
if lst_deep != lst_to_inspect:
return CountNestingLevels(lst_deep, cntr)
else:
return cntr if only_sublevels else cntr + 1
print(CountNestingLevels(dct_test)) # returns 6
print(CountNestingLevels(tpl_test)) # returns 2
Do you have any tips how to improve this code?