Chaining __getitem__ (d['a']['b']) not working with custom dict wrapper

Question

I have a dictionary where each key is always prepended with a particular string. The dict can look like this:

d = {
    "aa123":{
        "aa456": "456",
        "aa789": "789"
    }
}

So I am writing a wrapper where I can just query the dict without using the prepended string. Eg:

print(d["123"]["456"]) # --> should print "456"

So here's my wrapper:

class CustDict(dict):
    def __init__(self, *args, **kwargs):
        self.ns = kwargs.pop("namespace")
        super().__init__(*args, **kwargs)

    def __getitem__(self, key):
        key = f"{self.ns}{key}"
        return super().__getitem__(key)

When I use it I get the following error:

cust_d = CustDict(d, namespace="aa")
print(cust_d["123"]["456"])

I get the error:

KeyError: '456'

Now, I know this is happening because __getitem__ is returning an instance of dict instead of CustDict. But if I replace return super().__getitem__(key) with return CustDict(k, namespace=self.ns) I get other errors like ValueError: dictionary update sequence element #0 has length 1; 2 is required

Any solution for this would be appreciated.

Answer 1

First off, since you want to override the instance method of __getitem__ then you should not be subclassing from dict. If you inherit from dict then it will not even look at the instance method of __getitem__. You can learn more about that here. Instead use UserDict.

Some slight modifications to your code then make it look like the following:

from collections import UserDict

class CustDict(UserDict):
    def __init__(self, *args, **kwargs):
        self.ns = kwargs.pop("namespace")
        super().__init__(*args, **kwargs)

    def __getitem__(self, key):
        key = f"{self.ns}{key}"
        val = super().__getitem__(key)
        if isinstance(val, dict):
            return CustDict(val, namespace=self.ns)
        else:
            return val

cust_d = CustDict(d, namespace="aa")

cust_d["123"]
>> {'aa456': '456', 'aa789': '789'}

cust_d["123"]["456"]
>> '456'