python - ContextManager 访问调用 with block 的 locals()

Question

我正在尝试编写一个带有上下文管理器的多线程助手。这个想法是在一个 block 内定义一堆函数，上下文管理器“神奇地”负责调度和一切。简化的工作版本如下所示:

import contextlib

@contextlib.contextmanager
def multi_threaded(count):
    funcs = []
    yield funcs
    my_slice = int(count / len(funcs))
    for i, func in enumerate(funcs):
        start = my_slice * i
        func(start, start + my_slice)   


def spawn_many():
    dataset = [1, 2, 3, 4, 5]
    with multi_threaded(len(dataset)) as mt:
        def foo(start_idx, end):
            print("foo" + str(dataset[start_idx : end]))
        def bar(start_idx, end):
            print("bar" + str(dataset[start_idx : end]))
        mt.append(foo)
        mt.append(bar)

spawn_many()

这个示例有效，但我想删除这些行:

        mt.append(foo)
        mt.append(bar)

这样用户只需要定义函数而不需要将它们添加到集合中。为什么？因为它不太容易出错，而且我无法控制用这个库编写的代码。

问题是，在yield之后，我超出了发生def foo的范围，所以我不知道该范围内存在的locals() ，这基本上就是我需要知道其中定义了哪些函数的内容。有什么想法/技巧/鼓励的话吗？

感谢您的阅读!

Answer 1

装饰器可能会更好一点:

import contextlib

@contextlib.contextmanager
def multi_threaded(count):
    funcs = []
    yield funcs
    my_slice = int(count / len(funcs))
    for i, func in enumerate(funcs):
        start = my_slice * i
        func(start, start + my_slice)   

def add_to_flist(mt):
    def _add_to_flist(func):
        mt.append(func)
        return func
    return _add_to_flist

def spawn_many():
    dataset = [1, 2, 3, 4, 5]
    with multi_threaded(len(dataset)) as mt:
        @add_to_flist(mt)
        def foo(start_idx, end):
            print("foo" + str(dataset[start_idx : end]))
        @add_to_flist(mt)
        def bar(start_idx, end):
            print("bar" + str(dataset[start_idx : end]))

spawn_many()