4

Suppose I have a class that will spawn a thread and implements .__enter__ and .__exit__ so I can use it as such:

with SomeAsyncTask(params) as t:
    # do stuff with `t`
    t.thread.start()
    t.thread.join()

.__exit__ might perform certain actions for clean-up purposes (ie. removing temp files, etc.)

That works all fine until I have a list of SomeAsyncTasks that I would like to start all at once.

list_of_async_task_params = [params1, params2, ...]

How should I use with on the list? I'm hoping for something like this:

with [SomeAsyncTask(params) for params in list_of_async_task_params] as tasks:
    # do stuff with `tasks`
    for task in tasks:
        task.thread.start()
    for task in tasks:
        task.thread.join()
Derek 朕會功夫
  • 92,235
  • 44
  • 185
  • 247

3 Answers3

4

I think contextlib.ExitStack is exactly what you're looking for. It's a way of combining an indeterminate number of context managers into a single one safely (so that an exception while entering one context manager won't cause it to skip exiting the ones it's already entered successfully).

The example from the docs is pretty instructive:

with ExitStack() as stack:
    files = [stack.enter_context(open(fname)) for fname in filenames]
    # All opened files will automatically be closed at the end of
    # the with statement, even if attempts to open files later
    # in the list raise an exception

This can adapted to your "hoped for" code pretty easily:

import contextlib

with contextlib.ExitStack() as stack:
    tasks = [stack.enter_context(SomeAsyncTask(params))
             for params in list_of_async_task_params]
    for task in tasks:
        task.thread.start()
    for task in tasks:
        task.thread.join()
Blckknght
  • 100,903
  • 11
  • 120
  • 169
  • Awesome, but is there an equivalent method for Python2.7? It seems it was introduced in 3.3. – Derek 朕會功夫 Sep 07 '17 at 03:19
  • According to [the answer to question](https://stackoverflow.com/questions/34630393/python2-7-contextlib-exitstack-equivalent), there's a backport of the Python 3.5 version of the `contextlib` module for Python 2 named [`contextlib2`](https://contextlib2.readthedocs.io/en/stable/). @kichik's answer presents a rough start at a version you could create for yourself, though you might want to be more careful about `__enter__` methods throwing exceptions. – Blckknght Sep 07 '17 at 04:23
2

Note: Somehow I missed the fact that your Thread subclass was also a context manager itself—so the code below doesn't make that assumption. Nevertheless, it might be helpful when using more "generic" kinds of threads (where using something like contextlib.ExitStack wouldn't be an option).

Your question is a little light on details—so I made some up—however this might be close to what you want. It defines a AsyncTaskListContextManager class that has the necessary __enter__() and __exit__() methods required to support the context manager protocol (and associated with statements).

import threading
from time import sleep

class SomeAsyncTask(threading.Thread):
    def __init__(self, name, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.name = name
        self.status_lock = threading.Lock()
        self.running = False

    def run(self):
        with self.status_lock:
            self.running = True

        while True:
            with self.status_lock:
                if not self.running:
                    break
            print('task {!r} running'.format(self.name))
            sleep(.1)

        print('task {!r} stopped'.format(self.name))

    def stop(self):
        with self.status_lock:
            self.running = False


class AsyncTaskListContextManager:
    def __init__(self, params_list):
        self.threads = [SomeAsyncTask(params) for params in params_list]

    def __enter__(self):
        for thread in self.threads:
            thread.start()
        return self

    def __exit__(self, *args):
        for thread in self.threads:
            if thread.is_alive():
                thread.stop()
                thread.join()  # wait for it to terminate
        return None  # allows exceptions to be processed normally

params = ['Fee', 'Fie', 'Foe']
with AsyncTaskListContextManager(params) as task_list:
    for _ in range(5):
        sleep(1)
    print('leaving task list context')

print('end-of-script')

Output:

task 'Fee' running
task 'Fie' running
task 'Foe' running
task 'Foe' running
task 'Fee' running
task 'Fie' running
... etc
task 'Fie' running
task 'Fee' running
task 'Foe' running
leaving task list context
task 'Foe' stopped
task 'Fie' stopped
task 'Fee' stopped
end-of-script
martineau
  • 119,623
  • 25
  • 170
  • 301
  • `AsyncTaskListContextManager.__enter__()` seems to be missing `return self.threads`. – kichik Sep 07 '17 at 02:24
  • @kichik: Nice catch...fixed (but using a different return value). Thanks for pointing the potential issue out, although it wasn't really affecting the sample usage shown. – martineau Sep 07 '17 at 11:53
1

@martineau answer should work. Here's a more generic method that should work for other cases. Note that exceptions are not handled in __exit__(). If one __exit__() function fails, the rest won't be called. A generic solution would probably throw an aggregate exception and allow you to handle it. Another corner case is when you your second manager's __enter__() method throws an exception. The first manager's __exit__() will not be called in that case.

class list_context_manager:
  def __init__(self, managers):
    this.managers = managers

  def __enter__(self):
    for m in self.managers:
      m.__enter__()
    return self.managers

  def __exit__(self):
    for m in self.managers:
      m.__exit__()

It can then be used like in your question:

with list_context_manager([SomeAsyncTask(params) for params in list_of_async_task_params]) as tasks:
    # do stuff with `tasks`
    for task in tasks:
        task.thread.start()
    for task in tasks:
        task.thread.join()
kichik
  • 33,220
  • 7
  • 94
  • 114