import asyncio

from collections import UserDict
from typing import Any


class AwaitDict(UserDict):
    def __init__(self, dict=None, /, **kwargs):
        self._await_task: dict[str, asyncio.Future] = {}
        super().__init__(dict=None, **kwargs)

    def __setitem__(self, key: str, item: Any):
        if future := self._await_task.get(key):
            if future.done():
                raise Exception('Expected value has not been received yet')
            future.set_result(item)

        super().__setitem__(key, item)

    def __getitem__(self, item: str) -> Any:
        if item in self._await_task:
            self._await_task.pop(item)

        return super().__getitem__(item)

    async def await_elem(self, key: str) -> Any:
        future = asyncio.Future()
        self._await_task[key] = future
        await future
        return future.result()


async def set_elem(aw: AwaitDict, key: str, value: Any):
    await asyncio.sleep(3)
    aw[key] = value


async def main():
    my_dict = AwaitDict()
    asyncio.create_task(
        set_elem(
            aw=my_dict,
            key='123',
            value=123
        )
    )
    data = await my_dict.await_elem('123')
    print(data)
    print(my_dict['123'])

asyncio.run(main())