一个Python并发编程技巧:future当作字典的key当作中间值构建最终结果

代码如下

'''
我们现在想要得到一个字典: user_id: user_name
但是 子任务的方法只返回了user_id对应的user_name

1、同步构建结果比较容易，在for循环user_id的时候每次拿到结果后直接根据临时的user_id构建即可
2、如果是异步获取的话，这里有一个技巧，就是先构建一个 future 到 user_id 的字典,再在同步等待的时候反向构建最终的结果
'''
import time
import traceback
from functools import wraps
from concurrent.futures import ThreadPoolExecutor, as_completed


def timer(func):
    """
    简易函数执行时间装饰器，仅打印函数名和执行时间（秒）
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        execution_time = end_time - start_time

        # 仅打印函数名和执行时间（保留6位小数）
        print(f"{func.__name__} 执行时间: {execution_time:.6f} 秒")
        return result

    return wrapper


def get_user_name_by_id_delay_work(user_id: str) -> str:
    '''
    获取用户名的子任务
    '''

    # 模拟耗时
    time.sleep(0.2)

    return f'user_name-{user_id}'


@timer
def get_user_dict_sync(user_id_list: list[str]) -> dict:
    '''
    同步代码
    '''
    ret = {}
    if not user_id_list:
        return ret

    for uid in user_id_list:
        curr_user_name = get_user_name_by_id_delay_work(uid)
        ret[uid] = curr_user_name

    return ret


@timer
def get_user_dict_async(user_id_list: list[str]) -> dict:
    '''
    异步代码
    '''
    ret = {}
    if not user_id_list:
        return ret

    threading_pool = ThreadPoolExecutor(max_workers=10)

    # Notice 用于映射 future -> user_id，方便后续匹配结果
    future_to_uid = {}
    for uid in user_id_list:
        future = threading_pool.submit(
            get_user_name_by_id_delay_work,
            user_id=uid
        )
        # Notice 方便后面构建 user_id: user_name 的映射关系
        future_to_uid[future] = uid

    # 等待所有任务完成 收集结果
    for future in as_completed(future_to_uid):
        curr_uid = future_to_uid[future]
        try:
            # 获取子函数的用户名
            user_name = future.result()
            if user_name:
                ret[curr_uid] = user_name
        except Exception:
            print(f'并发查用户名失败! uid: {curr_uid}, 异常: {traceback.format_exc()}')

    return ret


if __name__ == '__main__':
    user_id_list = ['id1', 'id2', 'id3', 'id4', 'id5', 'id6', 'id7']

    ret_sync = get_user_dict_sync(user_id_list)
    print('ret_sync:>>> ', ret_sync)
    '''
    get_user_dict_sync 执行时间: 1.424339 秒
    ret_sync:>>>  {'id1': 'user_name-id1', 'id2': 'user_name-id2', 'id3': 'user_name-id3', 'id4': 'user_name-id4', 'id5': 'user_name-id5', 'id6': 'user_name-id6', 'id7': 'user_name-id7'}
    '''

    ret_async = get_user_dict_async(user_id_list)
    print('ret_async:>>> ', ret_async)
    '''
    get_user_dict_async 执行时间: 0.207139 秒
    ret_async:>>>  {'id3': 'user_name-id3', 'id1': 'user_name-id1', 'id2': 'user_name-id2', 'id4': 'user_name-id4', 'id6': 'user_name-id6', 'id7': 'user_name-id7', 'id5': 'user_name-id5'}
    '''

～～～