处理 KeyError 的上下文管理器 class

Question

我正在尝试实现一个 context manager class 来处理字典上的 KeyError。

想想这个：

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

有了这两个列表，我必须构建一个包含每个品牌和颜色销售额的二级字典，例如：

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

（对于此示例，仅将销售额作为随机数）。

我解决这个问题的实现最多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

上面函数的行为是预期的，但我想要一个用户定义的 class 来避免我们每次不得不面对这个问题时重复这段代码，我想一个上下文管理器 class 将是实现它的方法。

这是我所做的，但没有按预期工作：

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

这样我们就可以制作类似的东西：

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

但是 with 块内的迭代在上下文管理器首先处理 KeyError 之后停止，结果我只得到这个：{'Honda': {}}

Answer 1

您对 get_bikes_sales 的实现非常 pythonic (try/except)。
使用上下文管理器并不能解决问题 - 只是将它移到其他地方。

为什么不创建一个动态创建（任意嵌套）指令的函数：

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)

输出：

{'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}},
...

处理 KeyError 的上下文管理器 class

Context manager class which handles KeyError

python

python-3.x

exception

with-statement

keyerror