# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

import asyncio
import functools
from collections.abc import Callable, Iterable
from typing import Optional


def async_map(function: Callable, arguments: Iterable[Iterable], max_concurrent: Optional[int] = None) -> Iterable:
    """Asynchronously map the synchronous `function` over a list of
    `arguments` using asyncio."""

    if max_concurrent is not None and max_concurrent <= 0:
        raise ValueError('max_concurrent must be at least 1')

    function_list = [functools.partial(function, *x) for x in arguments]

    async def main():
        loop = asyncio.get_running_loop()
        if max_concurrent is None:
            tasks = [loop.run_in_executor(None, f) for f in function_list]
            return await asyncio.gather(*tasks)

        semaphore = asyncio.Semaphore(max_concurrent)

        async def run_with_semaphore(f):
            async with semaphore:
                return await loop.run_in_executor(None, f)

        tasks = [run_with_semaphore(f) for f in function_list]
        return await asyncio.gather(*tasks)

    return asyncio.run(main())
