import asyncio import inspect from typing import Any, Awaitable, Callable, Iterable, Iterator, List, Tuple, Type, Union from server.config import CONCURRENCY_LEVEL from server.utils import get_arg, set_arg from server.utils.parallel.results import DictResult, ListResult def iter_chunk(iterable: list, chunk_size: int = 1) -> Iterator: """Split collection of items on chunks. Args: iterable: Collection of items. chunk_size: Size of chunks. Returns: Iterator of chunks. """ count = len(iterable) for i in range(0, count, chunk_size): yield iterable[i : min(i + chunk_size, count)] # noqa async def execute_in_parallel( func: Callable, kwargs_list: List[dict], chunk_size: int = CONCURRENCY_LEVEL, result_type: Type[Union[ListResult, DictResult]] = ListResult, update_responses: Callable = None, ) -> list: """Execute function requests in parallel and sequential. Args: func: Function to call. kwargs_list: Kwargs list. chunk_size: A size of a chunk. result_type: Result type. update_responses: Update responses. Returns: Concatenated result. """ result = result_type.result_cls() for sequential_chunk in iter_chunk(kwargs_list, chunk_size): tasks = [] for kwargs in sequential_chunk: tasks.append(func(**kwargs)) responses = await asyncio.gather(*tasks) if update_responses: update_responses(responses, sequential_chunk) for response in responses: result = result_type.update(result, response) return result def request_in_chunks( chunk_size: int = None, chunks_count: int = CONCURRENCY_LEVEL, items_key: Union[str, Tuple[str]] = "ids", result_type: Type[Union[ListResult, DictResult]] = DictResult, update_func: Callable = None, ) -> Callable: """Execute API requests for chunks of items. Executes original request by synchronous chunks with inner asynchronous requests equal to chunk_count each of chunk_size size. Args: chunk_size: A size of a chunk. chunks_count: A number of chunks to run asynchronously. items_key: Args list key, it is possible to set multiple keys in tuple - this way the first not nullable-value key will be used. result_type: Wrapped function result type class. update_func: Optional custom function to merge results. Returns: Wrapped function. """ def wrapper(f: Callable): if isinstance(items_key, str): _items_key, _items_keys = items_key, (items_key,) elif isinstance(items_key, tuple): _items_key, _items_keys = None, items_key else: raise ValueError(f"request_in_chunks: items_key must be str or tuple, got {items_key} of {type(items_key)}") async def wrapped(*args, **kwargs): real_items_key = _items_key if not set(_items_keys) & kwargs.keys(): return await f(*args, **kwargs) if not real_items_key: # pick a first key with not empty value, or the first key in case if all values are empty real_items_key = next(filter(lambda k: bool(kwargs.get(k)), _items_keys), _items_keys[0]) items = kwargs[real_items_key] result = result_type.result_cls() _update_func = update_func or result_type.update sequential_chunk_size = chunks_count * (chunk_size or 1) for sequential_chunk in iter_chunk(items, sequential_chunk_size): if chunk_size: tasks = [] for concurrent_chunk in iter_chunk(sequential_chunk, chunk_size): current_kwargs = dict(kwargs) current_kwargs[real_items_key] = concurrent_chunk tasks.append(f(*args, **current_kwargs)) chunk_inner_results = await asyncio.gather(*tasks) for inner_result in chunk_inner_results: result = _update_func(result, inner_result) else: kwargs[real_items_key] = sequential_chunk chunk_result = await f(*args, **kwargs) result = _update_func(result, chunk_result) return result wrapped.__signature__ = inspect.signature(f) return wrapped return wrapper async def make_requests(requests: Iterable[Tuple[Callable, dict, Any, bool] or Tuple[Callable, dict]]) -> Iterable[Any]: """Make multiple requests in parallel and return response or default.json Args: requests: Requests definition: function, kwargs, default value, condition. Returns: Response or default. """ requests = [r if len(r) > 2 else (r[0], r[1], None, True) for r in requests] tasks = [i[0](**i[1]) for i in requests if i[3]] responses = await asyncio.gather(*tasks) result = [] response_index = 0 for index, request in enumerate(requests): if request[3]: result.append(responses[response_index]) response_index += 1 else: result.append(request[2]) return result async def _get_chunk( item_list: Iterable, field_name: str, f: Callable, args: tuple, kwargs: dict, args_spec: inspect.FullArgSpec, single_item_call: bool = False, ) -> tuple: """Process requests chunk. Args: item_list: Items list. field_name: Field name. f: Function to call. args: Function args. kwargs: Function kwargs. args_spec: Args spec. single_item_call: Call per single item. Returns: Chunk responses. """ tasks = [] for item in item_list: modified_args, modified_kwargs = set_arg( args, kwargs, args_spec, field_name, item if single_item_call else [item] ) tasks.append(f(*modified_args, **modified_kwargs)) return await asyncio.gather(*tasks) def request_per_item( field_name: str, default_value: list = None, sum_func: Callable[[Any, Any], Any] = None, sum_all_func: Callable = None, modify_items_func: Callable = None, single_item_call: bool = False, chunk_size: int = None, ): """Call function for each item of some list argument. Args: field_name: List arg field name. default_value: Default field value if None (for cases when None = ALL). sum_func: A function to sum up result's chunks. sum_all_func: A function to sum up all result's chunks at once. modify_items_func: A function to modify field value (ex: request related ISRC list). single_item_call: Call a function with single item, not a new list of this single item. chunk_size: Parallel chunk size. """ def inner(f: Callable): args_spec = inspect.getfullargspec(f) async def wrapped(*args, **kwargs) -> list: field_value = get_arg(args, kwargs, args_spec, field_name, from_default=False) or default_value if not field_value: return await f(*args, **kwargs) if not isinstance(field_value, (list, tuple, set)): field_value = [field_value] if modify_items_func: field_value = await modify_items_func(field_value, *args, **kwargs) if chunk_size: responses = [] for item_chunk in iter_chunk(field_value, chunk_size): responses_chunk = await _get_chunk( item_chunk, field_name, f, args, kwargs, args_spec, single_item_call ) responses.extend(responses_chunk) else: responses = await _get_chunk(field_value, field_name, f, args, kwargs, args_spec, single_item_call) if sum_all_func: result = sum_all_func(responses, args_spec, *args, **kwargs) else: result = responses[0] for chunk in responses[1:]: result = sum_func(result, chunk) if sum_func else result + chunk return result wrapped.__signature__ = inspect.signature(f) return wrapped return inner async def get_awaitable_or_default(f: Awaitable, condition: bool, default: Any, *args, **kwargs) -> Any: """ Check condition and return coroutine func or fake coroutine Args: f (): condition (Bool): default (Any): value to be returned in further call of coroutine obj *args: **kwargs: Returns: (Any) """ if condition: return await f(*args, **kwargs) return default