Skip to content

Add support for async functions in memo library #43

@petrhrobar

Description

@petrhrobar

Hi Vicent,

First, thank you for creating and maintaining the memo library. It has been incredibly useful for my projects. I am writing to request support for asynchronous functions in the memo library.

Here's an example of my use case:

import asyncio
from memo import memfile

@memfile("async_function.jsonl")
async def get_data():
    # Simulate some async work with asyncio.sleep
    await asyncio.sleep(1)
    return {"key1": "value1", "key2": "value2"}

# Example usage
async def main():
    result = await get_data()
    print(result)

# To run the example multiple times
if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    for _ in range(5):
        loop.run_until_complete(main())

and the error:

Traceback (most recent call last):
  File "/home/hrobar/msd_projects/dssi-nlp-gracs-chatbot/other/async_code.py", line 20, in <module>
    loop.run_until_complete(main())
  File "/home/hrobar/miniconda3/envs/gracs-llm/lib/python3.11/asyncio/base_events.py", line 654, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/home/hrobar/msd_projects/dssi-nlp-gracs-chatbot/other/async_code.py", line 13, in main
    result = await get_data()
                   ^^^^^^^^^^
  File "/home/hrobar/miniconda3/envs/gracs-llm/lib/python3.11/site-packages/memo/_base.py", line 108, in wrapper
    {**kwargs, **result},
    ^^^^^^^^^^^^^^^^^^^^
TypeError: 'coroutine' object is not a mapping
sys:1: RuntimeWarning: coroutine 'get_data' was never awaited
It appears that the memo library's decorator does not handle asynchronous functions properly, resulting in the TypeError: 'coroutine' object is not a mapping error.

I guess an additinal wrapper like this perhaps would be a way to go? (not sure how well it implements with yours code tho):

import asyncio
import orjson
import pathlib
from functools import wraps
from memo import memfile

def async_memfile(filepath: str, skip: bool = False):
    """
    Remembers input/output of an asynchronous function in a jsonl file on disk.

    Arguments:
        filepath: path to write data to
        skip: skips the calculation if kwargs appear in data already
    """

    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            result = await func(*args, **kwargs)
            if skip:
                if pathlib.Path(filepath).exists():
                    with open(filepath, "r") as f:
                        datalist = [orjson.loads(line) for line in list(f)]
                else:
                    datalist = []
            with open(filepath, "a") as f:
                if skip and _contains(kwargs, datalist):
                    return None
                ser = orjson.dumps(
                    {**kwargs, **result},
                    option=orjson.OPT_NAIVE_UTC | orjson.OPT_SERIALIZE_NUMPY,
                )
                f.write(ser.decode("utf-8") + "\n")
            return result

        return wrapper

    return decorator

Thank you for considering this feature request. I am looking forward to your response.

Best regards,

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions