74 lines
2.6 KiB
Python
74 lines
2.6 KiB
Python
import asyncio
|
|
import inspect
|
|
import multiprocessing as mp
|
|
import queue
|
|
import threading
|
|
import traceback
|
|
from multiprocessing.synchronize import Event
|
|
from typing import Any, Awaitable, Callable, List, TypeVar, Union, cast
|
|
|
|
import dill
|
|
|
|
from .map_result import MapResult
|
|
|
|
T = TypeVar("T")
|
|
V = TypeVar("V")
|
|
|
|
|
|
def mapper_function(
|
|
input_queue: Union[mp.Queue, queue.Queue],
|
|
output_queue: Union[mp.Queue, queue.Queue],
|
|
should_stop: Union[Event, threading.Event],
|
|
func: Union[bytes, Callable[[T], V], Callable[[T], Awaitable[V]]],
|
|
) -> None:
|
|
try:
|
|
if isinstance(func, bytes):
|
|
func = cast(Callable[[T], V], dill.loads(func))
|
|
|
|
is_asynchronous = inspect.iscoroutinefunction(func)
|
|
|
|
last_chunk: List[MapResult] = []
|
|
while not should_stop.wait(0.1):
|
|
if not last_chunk:
|
|
try:
|
|
input_chunk = input_queue.get_nowait()
|
|
if is_asynchronous:
|
|
|
|
async def safe(i: int, value: T) -> Any:
|
|
try:
|
|
return MapResult(
|
|
i,
|
|
await cast(Callable[[T], Awaitable[V]], func)(
|
|
value
|
|
),
|
|
)
|
|
except Exception as e:
|
|
return MapResult(i, None, e, traceback.format_exc())
|
|
|
|
async def main() -> List[MapResult]:
|
|
return await asyncio.gather(
|
|
*[safe(i, v) for i, v in input_chunk]
|
|
)
|
|
|
|
last_chunk = asyncio.run(main())
|
|
else:
|
|
for i, value in input_chunk:
|
|
try:
|
|
last_chunk.append(MapResult(i, func(value)))
|
|
except Exception as e:
|
|
last_chunk.append(
|
|
# `e` has to be stringified in order to avoid any
|
|
# surprising serialization error when returning it
|
|
MapResult(i, None, str(e), traceback.format_exc())
|
|
)
|
|
except queue.Empty:
|
|
pass
|
|
|
|
if last_chunk:
|
|
try:
|
|
output_queue.put_nowait(last_chunk)
|
|
last_chunk = []
|
|
except queue.Full:
|
|
pass
|
|
except (KeyboardInterrupt, BrokenPipeError):
|
|
pass
|