使用ContextManager块访问调用的本地变量()

Question

我正在尝试编写一个具有上下文管理器的多线程助手。这个想法是在一个块中定义一组函数，上下文管理器“神奇地”负责调度和所有事情。一个简化的工作版本如下所示：

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)

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

问题是，在输出之后，我超出了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()

Answer 2

我读过那个this is not possible，至少不是没有丑陋的黑客，但我认为我的解决方案最终并没有那么丑陋：

您在创建时将locals()字典传递给contextmanager，contextmanager在返回后询问该字典，以收集任何可调用的数据：

@contextlib.contextmanager
def multi_threaded(block_locals, count):
    yield

    funcs = [fn for fn in block_locals.values() if callable(fn)]

    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(locals(), len(dataset)):
        def foo(start_idx, end):
            print("foo" + str(dataset[start_idx : end]))
        def bar(start_idx, end):
            print("bar" + str(dataset[start_idx : end]))

        # Re-sync locals-dict handed earlier to multi_threaded().
        locals()

spawn_many()

请注意，这个技巧之所以有效，是因为块中最后一个对locals()的调用。似乎Python仅在调用locals()时才同步locals()-dictionary <-->函数局部变量。如果没有最后一次调用，multi_threaded会将{'dataset': [1, 2, 3, 4, 5]}视为本地人。