Python Multiprocessing Map Dictionary

README.md

Map Dictionary Using Python Multiprocessing

Simple utility function that can transform the values in a dictionary into a new dictionary using multiple threads, if used with Python Multiprocessing.

Installation

Copy the content of the attached dict_pool_map.py-file. A non-typed version is also attached to showcase an easy-to-read implementation (dict_pool_map_no_typing.py).

Usage

import multiprocessing
import time


def double_it(x: int) -> int:
    time.sleep(1)
    return x * 2


inp = {1: 1, 2: 2, 3: 3}

with multiprocessing.Pool() as p:
    out = dict_pool_map(p, double_it, inp)

print(out)  # {1: 2, 2: 4, 3: 6}

NOTE: For non-concurrent workloads, use a dictionary comprehension instead.

out = {key: double_it(value) for key, value in inp.items()}

dict_pool_map.py

from typing import Callable, Protocol, TypeVar, cast

Input = TypeVar("Input")
Output = TypeVar("Output")
Key = TypeVar("Key")


class PoolLike(Protocol):
    def map(
        self,
        func: Callable[[Input], Output],
        iterable: list[Input],
    ) -> list[Output]: ...


def dict_pool_map(
    pool: PoolLike,
    callback: Callable[[Input], Output],
    x: dict[Key, Input],
):
    keys, values = cast(tuple[list[Key], list[Input]], zip(*x.items()))

    results = pool.map(callback, values)

    return {key: res for key, res in zip(keys, results)}

dict_pool_map_no_typing.py

def dict_pool_map(pool, callback, x):
    keys, values = zip(*x.items())
    results = pool.map(callback, values)
    return {key: res for key, res in zip(keys, results)}