Method map

src/memos/context/context.py:279–313

Returns an iterator equivalent to map(fn, iter). Automatically propagates the current thread's context to worker threads.

(
        self,
        fn: Callable[..., T],
        *iterables: Any,
        timeout: float | None = None,
        chunksize: int = 1,
    )

        return super().submit(wrapper, *args, **kwargs)

    def map(
        self,
        fn: Callable[..., T],
        *iterables: Any,
        timeout: float | None = None,
        chunksize: int = 1,
    ) -> Any:
        """
        Returns an iterator equivalent to map(fn, iter).
        Automatically propagates the current thread's context to worker threads.
        """
        main_trace_id = get_current_trace_id()
        main_api_path = get_current_api_path()
        main_env = get_current_env()
        main_user_type = get_current_user_type()
        main_user_name = get_current_user_name()
        main_context = get_current_context()

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            if main_context:
                # Create and set new context in worker thread
                child_context = RequestContext(
                    trace_id=main_trace_id,
                    api_path=main_api_path,
                    env=main_env,
                    user_type=main_user_type,
                    user_name=main_user_name,
                )
                child_context._data = main_context._data.copy()
                set_request_context(child_context)

            return fn(*args, **kwargs)

        return super().map(wrapper, *iterables, timeout=timeout, chunksize=chunksize)


# Type for trace_id getter function