Tremmillicious
Tremmillicious

Reputation: 506

How to get the depth items of dict in python and discard the rest

I have a complicated json file coming from somewhere and I am trying to get n level dictionary and discard the rest. (so this is not flattening as such). For example given a dict of countries and their subcountries and cities

my_dict = {
  'USA':['California', 'Arizona', {'Ohio':'Columbus'}],
  'Zimbabwe':{'Manicaland':['Mutare', 'Nyanga']},
  'New Zealand':''
}

Calling for level 1 to get list of countries

n_level(my_dict, 1)

Should give

['USA', 'Zimbabwe', 'New Zealand']

#calling for level 2
n_level(my_dict,  2)

#should give 

[
  {'USA':['California', 'Arizona', 'Ohio']},
  {'Zimbabwe':'Manicaland'}, #<--str
  'New Zealand'
]

I have tried this but not working no idea

my_dict = {
  'USA':['California', 'Arizona', {'Ohio':'Columbus'}],
  'Zimbabwe':{'Manicaland':['Mutare', 'Nyanga']}
}
def make_leaf_list(d):
   
    r = []
    if isinstance(d, collections.abc.Mapping):
        for k, v in d.items():
            r.append(k)
    elif isinstance(d, list):
        for i in d:
            r.extend(make_leaf_list(i))
    else:
        return r.append(d)
    
    return r

def n_level(o, depth):
    
    if depth == 0:
            return make_leaf_list(o)
   
    if isinstance(o, collections.abc.Mapping):        
        r = {}
        for k, v in o.items():
            r[k] = n_level(v, depth-1)
        return r
    elif isinstance(o, list):
        r = []
        for i in o:
            r.append(update(i, depth-1))
        return r
    else:
        return o

#and calling it
n_level(my_dict, 3)

Any idea for the approach? Thanks

Upvotes: 3

Views: 459

Answers (2)

Ajax1234
Ajax1234

Reputation: 71461

You can use a recursive generator function:

my_dict = {'USA': ['California', 'Arizona', {'Ohio': 'Columbus'}], 'Zimbabwe': {'Manicaland': ['Mutare', 'Nyanga']}, 'New Zealand': ''}
def to_level(d, l):
   if l: #make sure level is not 0
      if not isinstance(d, (list, dict)):
         yield d #not a dictionary or list, does not need to be traversed
      elif isinstance(d, dict):
         for a, b in d.items(): #element is a dictionary, have to transform it
            if not (n:=list(to_level(b, l-1))): #check this is the last level
               yield a #current level is `1`, so only yield back the key of the dictionary (no need to traverse the values, as we are out of depth)
            else: #at a level `l > 1`, so we need to transform the key's value
               if isinstance(b, (str, list)) or (len(b) == 1 and not isinstance(b, list)):
                  n = n[0] #original value `b` was a dictionary with a single key, a string, or a list
               elif isinstance(b, dict) and all(isinstance(i, dict) for i in n):
                  n = dict([j for k in n for j in k.items()]) #`b` was a dictionary with more than one key, or a dictionary with multiple additional levels to traverse
               yield {a:n} #yield back the key and its transformed value as its own dictionary
      else:
         yield [j for k in d for j in to_level(k, l)] #value to transform is a list, simply map `to_level` to each element and yield the resulting list

def n_level(d, l):
   return list(to_level(d, l))

result = n_level(my_dict, 1)
result1 = n_level(my_dict, 2)
result2 = n_level(my_dict, 3)

Output:

['USA', 'Zimbabwe', 'New Zealand']
[{'USA': ['California', 'Arizona', 'Ohio']}, {'Zimbabwe': 'Manicaland'}, {'New Zealand': ''}]
[{'USA': ['California', 'Arizona', {'Ohio': 'Columbus'}]}, {'Zimbabwe': {'Manicaland': ['Mutare', 'Nyanga']}}, {'New Zealand': ''}]

This will also work on expanded dictionarys with subdictionaries containing more than one key:

my_dict = {'USA': ['California', 'Arizona', 
                {'Ohio': 'Columbus', 'New Hampshire':'Concord'}], 
          'Zimbabwe': {'Manicaland': ['Mutare', 'Nyanga'], 
                       'Russia':'Moscow'}, 
          'New Zealand': ''}
for i in range(1, 4):
   print(n_level(my_dict, i))

Output:

['USA', 'Zimbabwe', 'New Zealand']
[{'USA': ['California', 'Arizona', 'Ohio', 'New Hampshire']}, {'Zimbabwe': ['Manicaland', 'Russia']}, {'New Zealand': ''}]
[{'USA': ['California', 'Arizona', {'Ohio': 'Columbus'}, {'New Hampshire': 'Concord'}]}, {'Zimbabwe': {'Manicaland': ['Mutare', 'Nyanga'], 'Russia': 'Moscow'}}, {'New Zealand': ''}]

Upvotes: 4

ab.fe
ab.fe

Reputation: 119

Here is a solution from where you can start, it is working with level 1 and 2, I don't know how many levels do you have, So this used type of element: list, dict or str and according to the type it knows if it had more depth or if it is the last element:

def n_level(my_dict, depth):
    if depth == 1:
            return list(my_dict.keys())
    if depth == 2:
      result={}
      for i in list(my_dict.keys()):
        # print(my_dict[i])
        if type(my_dict[i])==dict:
          # print(list(my_dict[i].keys())[0])
          add_to_dict(result,list(my_dict[i].keys())[0],i)

        elif type(my_dict[i])==list:
          for leaf in my_dict[i]:
            if type(leaf)==str:
              # print(leaf)
              add_to_dict(result,leaf,i)
            elif type(leaf)==dict:
              # print(list(leaf.keys())[0])
              add_to_dict(result,list(leaf.keys())[0],i)
        elif my_dict[i]=='':
          add_to_dict(result,my_dict[i],i)
      return result

When tested with level 2 :

n_level(my_dict,2)

it gaves:

{'New Zealand': '',
 'USA': ['California', 'Arizona', 'Ohio'],
 'Zimbabwe': 'Manicaland'}

and with level 1 :

n_level(my_dict,1)
['USA', 'Zimbabwe', 'New Zealand']

Upvotes: 3

Related Questions