Context manager class which handles KeyError

Question

I'm trying to implement a context manager class which handles KeyError on dictionaries.

Think of this:

bikes = ['Honda', 'Yamaha', 'Kawasaki', 'Suzuki']
colors = ['Red', 'Blue', 'Green', 'White', 'Black']

Having these two lists, I have to build a two-evels dictionary with sales per brand and color, e.g:

bike_sales_by_color = {
    'Honda': {
        'Red': 100,
        'Blue': 125,
    },
    'Yamaha': {
        'Black': 50,
        'White': 60,
    },
    # etc...
}

(For this example take the sales amounts as just random numbers).

My implementation to solve this problem was the most ordinary/regular:

def get_bikes_sales():
    bikes_sales = {}
    for bike in bikes:  # iterate over two lists
        for color in colors:
            try:  # try to assign a value to the second-level dict.
                bikes_sales[bike][color] = 100  # random sales int value
            except KeyError:  # handle key error if bike is not yet in the first-level dict.
                bikes_sales[bike] = {}
                bikes_sales[bike][color] = 100  # random sales int value
    return bikes_sales

The behavior of the function above is the expected, but I would like a user-defined class to save us to repeat this code every time we have to face this problem and I thought a context manager class would be the way to achieve it.

This is what I did but it doesn't work as intended:

class DynamicDict:
    """
    Context manager class that immediately creates a new key with the intended value
    and an empty dict as the new created key's value if KeyError is raised.
    Useful when trying to build two or more level dictionaries.
    """

    def __init__(self):
        self.dictionary = {}

    def __enter__(self):
        return self.dictionary

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is KeyError:
            self.dictionary[exc_val.args[0]] = {}
            return True

So that we can just make something like:

with DynamicDict() as bikes_sales:
    for bike in bikes:
        for color in colors:
            bikes_sales[bike][color] = 100

But the iteration inside the with block stops after first KeyError handled by the context manager and I get as result only this: {'Honda': {}}

Answer 1

Your implementation of get_bikes_sales is quite pythonic (try/except).
Using a context manager doesn't solve the problem - just moves it somewhere else.

Why not creating a function that creates (arbitrary nested) dicts dynamically:

import itertools
import pprint


def generate_nested_dict(*categories, values):
    result = {}

    # create all combinations of all categories, alternatively recursion could be used.
    for tree in (x for x in itertools.product(*categories)):
        _dictHandle = result  # keeps track of the parent level with in the dict

        # Each tree is Honda->Red, Honda->Blue, ...
        for i, k in enumerate(tree):
            if k not in _dictHandle:
                if i < len(tree) - 1:
                    # add nested dict level
                    _dictHandle[k] = {}
                else:
                    # nested level exists
                    if len(values) == 1:
                        _dictHandle[k] = values[0]
                    else:
                        _dictHandle[k] = values.pop(0)
                    # add value
            _dictHandle = _dictHandle[k]
    return result


bikes = ['Honda', 'Yamaha', 'Kawasaki', 'Suzuki']
fuels = ['Petrol', 'Diesel', 'Electric', 'Soda']
colors = ['Red', 'Blue', 'Green', 'White', 'Black']
sales = [
    (100 * (i + 1)) + (10 * (j + 1)) for i in range(len(bikes))
    for j in range(len(colors))
]

# different values
bike_sales_by_color = generate_nested_dict(bikes, colors, values=sales)
pprint.pprint(bike_sales_by_color)

# different values and one category more
bike_sales_by_fuel_and_color = generate_nested_dict(
    bikes, fuels, colors, values=[100]
)
pprint.pprint(bike_sales_by_fuel_and_color)

Out:

{'Honda': {'Black': 150, 'Blue': 120, 'Green': 130, 'Red': 110, 'White': 140},
 'Kawasaki': {'Black': 350,
              'Blue': 320,
              'Green': 330,
              'Red': 310,
              'White': 340},
...
{'Honda': {'Diesel': {'Black': 100,
                      'Blue': 100,
                      'Green': 100,
                      'Red': 100,
                      'White': 100},
           'Electric': {'Black': 100,
                        'Blue': 100,
                        'Green': 100,
                        'Red': 100,
                        'White': 100},
           'Petrol': {'Black': 100,
                      'Blue': 100,
                      'Green': 100,
                      'Red': 100,
                      'White': 100},
           'Soda': {'Black': 100,
                    'Blue': 100,
                    'Green': 100,
                    'Red': 100,
                    'White': 100}},
...