From 8306a9f3ff4c231b145e6d2a2d264ce7cfe4973a Mon Sep 17 00:00:00 2001 From: Dmitriy Zayakin Date: Thu, 1 Aug 2024 16:32:41 +0300 Subject: [PATCH] [#276] Context manager for parralel func Signed-off-by: Dmitriy Zayakin --- src/frostfs_testlib/testing/parallel.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/frostfs_testlib/testing/parallel.py b/src/frostfs_testlib/testing/parallel.py index 9c36118..0549e61 100644 --- a/src/frostfs_testlib/testing/parallel.py +++ b/src/frostfs_testlib/testing/parallel.py @@ -1,10 +1,22 @@ import itertools from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager from typing import Callable, Collection, Optional, Union MAX_WORKERS = 50 +@contextmanager +def parallel_workers_limit(workers_count: int): + global MAX_WORKERS + original_value = MAX_WORKERS + MAX_WORKERS = workers_count + try: + yield + finally: + MAX_WORKERS = original_value + + def parallel( fn: Union[Callable, list[Callable]], parallel_items: Optional[Collection] = None,