Return __getitem__ method overridden object-CodePudding

The requirement is to raise a custom exception when KeyError happens. How do I make sure that the returned object via __getitem__() also has similar behavior?

  1 class AirflowCfg(dict):
  2     def __getitem__(self, key):
  3         """
  4         This ensures that if the config isn't provided,
  5         it's an AirflowException
  6         """
  7         try:
  8             return super().__getitem__(key)
  9         except KeyError:
 10             raise AirflowException(f"{key} was not provided")
 11
 12 x = AirflowCfg({"x": {"a": 1}})
 13 print(x["x"])
 14 print(x["x"]["a"])
 15 print(x["x"]["b"])

The above definitely doesn't make sense and it doesn't work, output:-

➤ python3 foo.py
{'a': 1}
1
Traceback (most recent call last):
  File "/Users/foo/foo.py", line 15, in <module>
    print(x["x"]["b"])
KeyError: 'b'

I've this which kinda of works, but I'm sure Python has a better way.

class ConfigDict(dict):
    def __getitem__(self, key):
        """
        This ensures that if the config isn't provided,
        it's an AirflowException
        """
        try:
            return super().__getitem__(key)
        except KeyError:
            raise AirflowException(f"{key} was not provided") from None

class AirflowCfg(dict):
    def __getitem__(self, key):
        try:
            return ConfigDict(super().__getitem__(key))
        except KeyError:
            pass

CodePudding user response：

I guess, I just wasn't thinking straight. After reading comments, It should be as easy as this:-

  1 class AirflowCfg(dict):
  2     def __getitem__(self, key):
  3         try:
  4             if isinstance(super().__getitem__(key), dict):
  5                 return AirflowCfg(super().__getitem__(key))
  6             else:
  7                 return super().__getitem__(key)
  8         except KeyError:
  9             raise AirflowException(f"{key} was not provided") from None

CodePudding user response：

It's not much better, but you can simplify the dict wrapper returned to avoid overhead in the case where the key is found:

class ConfigDict(dict):
    def __missing__(self, key):  # Called by dict.__getitem__ when the key is missing
                                 # Normally used to provide a default value, but raising
                                 # your own exception works too
        raise AirflowException(f"{key} was not provided")

class AirflowCfg(dict):
    def __getitem__(self, key):
        try:
            return ConfigDict(super().__getitem__(key))
        except KeyError:
            raise AirflowException(f"{key} was not provided")

This does shallow copy the contents of the dict into the new ConfigDict. This can impose some non-trivial costs, so you may want to write a slightly more complex implementation that contains a dict rather than inheriting from a dict, so it forms a view of the underlying dict with just slightly different behavior, without requiring a copy operation:

from collections.abc import MutableMapping

class ConfigDict(MutableMapping):
    __slots__ = 'dct'
    def __init__(self, dct):
        self.dct = dct
    def __getitem__(self, key):
        try:
            return self.dct[key]
        except KeyError:
            raise AirflowException(f"{key} was not provided") from None
    def __setitem__(self, key, value):
        self.dct[key] = value
    def __delitem__(self, key):
        del self.dct[key]  # Can wrap as in getitem case if desired, but you might need to implement some other methods like get, since it's provided with the assumption a missing item will raise KeyError
    def __iter__(self):
        return iter(self.dct)
    def __len__(self):
        return len(self.dct)