With Clause for Multiprocessing in Python

user_2208112 0
Viewed 11661

In python 3, you can now open a file safely using the with clause like this:

with open("stuff.txt") as f:
    data = f.read()

Using this method, I don't need to worry about closing the connection

I was wondering if I could do the same for the multiprocessing. For example, my current code looks like: 

pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
pool.starmap(function,list)
pool.close()
pool.join()

Is there any way I could use a with clause to simplify this?

with-statement python-3.x multiprocessing python
2 Answers

Although its more than what the OP asked, if you want something that will work for both Python 2 and Python 3, you can use:

# For python 2/3 compatibility, define pool context manager
# to support the 'with' statement in Python 2
if sys.version_info[0] == 2:
    from contextlib import contextmanager
    @contextmanager
    def multiprocessing_context(*args, **kwargs):
        pool = multiprocessing.Pool(*args, **kwargs)
        yield pool
        pool.terminate()
else:
    multiprocessing_context = multiprocessing.Pool

After that, you can use multiprocessing the regular Python 3 way, regardless of which version of Python you are using. For example:

def _function_to_run_for_each(x):
       return x.lower()
with multiprocessing_context(processes=3) as pool:
    results = pool.map(_function_to_run_for_each, ['Bob', 'Sue', 'Tim'])   
print(results)

Now, this will work in Python 2 or Python 3.

user_36387680
Related