PeacefulBY
PeacefulBY

Reputation: 163

How to get values with any depth in dict via *args parameter? (python)

For example, here I have:

doc = { 'A':1, 'B':1, 'C':{'C-A':2, 'C-B':{'C-B-A':3}}

then I define a func like this:

def get(doc, *args):
    if(len(args) == 1):
        return doc[args[0]]
    if(len(args) == 2):
        return doc[args[0]][args[1]]
    if(len(args) == 3):
        return doc[args[0]][args[1]][args[2]]

so I can get values like this:

get(doc,'A') //return 1
get(doc,'C','C-A') //return 2
get(doc,'C','C-B') //return {'C-B-A':3}
get(doc,'C','C-B','C-B-A') //return 3

Now my question is, if doc has any depth, how to rewrite func get?

Upvotes: 0

Views: 190

Answers (2)

Alfe
Alfe

Reputation: 59516

Use recursion:

def get(doc, *args):
    return get(doc[args[0]], *args[1:]) if args else doc

Or, if for whatever reason you prefer doing it iteratively

def get(doc, *args):
   for arg in args:
      doc = doc[arg]
   return doc

And that's damn near to @Bakuriu's answer.

Upvotes: 2

Bakuriu
Bakuriu

Reputation: 101989

You could use reduce:

In [17]: doc = { 'A':1, 'B':1, 'C':{'C-A':2, 'C-B':{'C-B-A':3}}}
    ...: 
    ...: def get(doc, *args):
    ...:     return reduce(dict.get, args, doc)

In [18]: get(doc, 'A')
Out[18]: 1

In [19]: get(doc, 'C', 'C-A')
Out[19]: 2

In [20]: get(doc, 'C', 'C-B', 'C-B-A')
Out[20]: 3

You can think of reduce as if it was:

def reduce(function, arguments, initializer):
    a = initializer
    for b in arguments:
        a = function(a, b)
    return a

I.e. it "accumulates" values calling a binary function over the given arguments. In your case first a is doc. Then, it becomes the result of doc[arguments[0]]. If arguments was length one the iteration stops and the result is returned, otherwise it is one of the sub-dicts and at the next iteration the loop will take a value from that sub-dict.


Note that using recursion means limiting to a depth of sys.getrecursionlimit(), which usually is about 1000, and it will be much slower since you are doing at least two times the function calls to retrieve the value:

In [84]: def get(doc, *args):
    ...:     return reduce(dict.get, args, doc)

In [85]: def get2(doc, *args):
    ...:     return get2(doc[args[0]], *args[1:]) if args else doc

In [86]: %timeit get(doc, 'C', 'C-B', 'C-B-A')
1000000 loops, best of 3: 621 ns per loop

In [87]: %timeit get2(doc, 'C', 'C-B', 'C-B-A')
1000000 loops, best of 3: 1.04 us per loop

In [88]: d = make_dict(depth=350) 

In [89]: %timeit get(d, *range(350))
10000 loops, best of 3: 38.9 us per loop

In [90]: %timeit get2(d, *range(350))
1000 loops, best of 3: 973 us per loop

And this using a shallow dict. If you try with a dict with bigger depths the recursive solution will get slower and slower compared to the functional solution.

Upvotes: 2

Related Questions