| # Copyright 2021 The Chromium Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| from __future__ import print_function |
| |
| import contextlib |
| from typing import Generator, Optional |
| |
| from pathos import pools |
| |
| |
| def GetProcessPool(nodes: Optional[int] = None) -> pools.ProcessPool: |
| """Returns a pathos.pools.ProcessPool instance. |
| |
| Split out for ease of unittesting since pathos can still run into pickling |
| issues with MagicMocks used in tests. |
| |
| Args: |
| nodes: How many processes processes to spawn in the process pool. |
| |
| Returns: |
| A pathos.pools.ProcessPool instance. |
| """ |
| return pools.ProcessPool(nodes=nodes) |
| |
| |
| @contextlib.contextmanager |
| def GetProcessPoolContext( |
| nodes: Optional[int] = None) -> Generator[pools.ProcessPool, None, None]: |
| try: |
| pool = GetProcessPool(nodes) |
| yield pool |
| finally: |
| pool.close() |
| pool.join() |
