Source code for swh.osv.utils

# Copyright (C) 2025-2026  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU Affero General Public License version 3, or any later version
# See top-level LICENSE file for more information

import asyncio
from pathlib import Path
import time
from typing import (
    Any,
    AsyncIterable,
    AsyncIterator,
    Awaitable,
    Callable,
    Generic,
    Iterable,
    TypeVar,
    Union,
)




def format_number(n: int | float) -> str:
    if n < 10:
        return str(n)
    elif n < 2_000:
        return str(int(n))
    elif n < 2_000_000:
        return f"{int(n / 1_000)}k"
    else:
        return f"{int(n / 1_000_000)}M"



_SENTINEL = object()
"""Object added at the end of a queue to tell workers to stop.

Alternative to ``asyncio.QueueShutDown`` that does not require Python >= 3.13"""

T = TypeVar("T")
U = TypeVar("U")




async def iter_queue(queue: asyncio.Queue[T]) -> AsyncIterator[T]:
    """Turns an async queue into an async iterator"""
    while True:
        item = await queue.get()
        if item is _SENTINEL:
            return
        yield item





async def concurrent_async_map(
    f: Callable[[T], Awaitable[U]],
    it: Union[Iterable[T], AsyncIterable[T]],
    max_concurrency: int = 50,
    buffer: int = 1000,
) -> AsyncIterator[U]:
    """applies ``f`` to each item in ``it``, with up to ``max_concurrency`` calls
    to ``f`` running at the same time."""
    work: asyncio.Queue[Any] = asyncio.Queue(maxsize=buffer)
    results: asyncio.Queue[Any] = asyncio.Queue(maxsize=buffer)
    remaining_workers = max_concurrency

    async def worker() -> None:
        nonlocal remaining_workers
        async for work_item in iter_queue(work):
            await results.put(await f(work_item))
        remaining_workers -= 1
        if remaining_workers == 0:
            await results.put(_SENTINEL)

    async def feeder():
        if hasattr(it, "__aiter__"):
            async for work_item in it:
                await work.put(work_item)
        else:
            for work_item in it:
                await work.put(work_item)
        for _ in range(max_concurrency):
            await work.put(_SENTINEL)

    async with asyncio.TaskGroup() as tg:
        tg.create_task(feeder())
        for _ in range(max_concurrency):
            tg.create_task(worker())
        async for result in iter_queue(results):
            yield result



D = TypeVar("D", bound=Union[dict, set])




class Cache(Generic[D]):
    def __init__(
        self,
        path: Path,
        initial_value: D,
        serializer: Callable[[D], str],
        deserializer: Callable[[str], D],
    ):
        self._data = initial_value
        self._path = path
        self._lock = asyncio.Lock()
        self._serializer = serializer
        self._deserializer = deserializer
        self._last_save_ts = 0.0



    async def load(self) -> None:
        async with self._lock:
            if self._path.exists():
                self._data.update(self._deserializer(self._path.read_text()))


    def __contains__(self, key) -> bool:
        return key in self._data

    def __getitem__(self, key):
        return self._data[key]

    def __getattr__(self, name):
        return getattr(self._data, name)

    def __setitem__(self, key, value):
        self._data[key] = value



    async def save(self) -> None:
        now = time.time()
        if self._last_save_ts + 10.0 > now:
            # don't bother saving less than 10s of changes
            return
        self._last_save_ts = now
        data = self._serializer(self._data)
        async with self._lock:
            self._path.write_text(data)