An executor that runs tasks on the main thread.

main_thread_executor.py

"""An executor that runs tasks on the main thread.

This can come in handy for debugging. For example, enables you to inspect and
set breakpoints on the underlying task, which is not easily done when using
separate threads.

Usually one would use this executor when the user provided `--n-jobs=0`.

"""
import logging
import concurrent.futures


class MainThreadExecutor(concurrent.futures.Executor):
    """An executor that runs tasks on the main thread."""

    def __init__(self, *args, **kwargs):
        logging.warning(f"MainThreadExecutor:: Ignoring args={args} and kwargs={kwargs}")
        self._shutdown = False

    def submit(self, fn, *args, **kwargs):
        if self._shutdown:
            raise RuntimeError("Cannot schedule new futures after shutdown")

        future = concurrent.futures.Future()
        try:
            result = fn(*args, **kwargs)
            future.set_result(result)
        except Exception as e:
            future.set_exception(e)

        return future

    def map(self, fn, *iterables, timeout=None, chunksize=1):
        return [
            self.submit(fn, *args).result()
            for args in zip(*iterables)
        ]

    def shutdown(self, wait=True):
        self._shutdown = True


if __name__ == "__main__":

    # Example usage:
    import threading
    def example_task(name):
        print(f"Executing task {name} on thread {threading.current_thread().name}")

    with MainThreadExecutor() as executor:
        results = list(executor.map(example_task, range(5)))