Implement AsyncSonyFlake

Author: ZipFile

README.rst

@@ -45,6 +45,23 @@ Turbo mode:
         t = timeit(lambda: [next(sf) for _ in range(1000)], number=1000)
         print(f"Speed: 1M ids / {t:.2f}sec with {count} machine IDs")
 
+Async:
+
+.. code-block:: python
+
+    from anyio import sleep  # AsyncSonyFlake supports both asyncio and trio
+    from sonyflake_turbo import AsyncSonyFlake, SonyFlake
+
+    sf = SonyFlake(0x1337, 0xCAFE, start_time=1749081600)
+    asf = AsyncSonyFlake(sf, sleep)
+
+    print(await asf)
+    print(await asf(5))
+
+    async for id_ in asf:
+        print(id_)
+        break
+
 Important Notes
 ===============
 

pyproject.toml

@@ -31,6 +31,8 @@ dev = [
     "isort",
     "mypy",
     "pytest",
+    "pytest-asyncio",
+    "pytest-trio",
 ]
 
 [tool.setuptools]
@@ -66,4 +68,4 @@ warn_unused_configs = true
 
 [tool.cibuildwheel]
 test-command = "pytest --import-mode=importlib -p no:cacheprovider {project}/tests"
-test-requires = ["pytest"]
+test-requires = ["pytest", "pytest-asyncio", "pytest-trio"]

setup.py

@@ -34,6 +34,7 @@
                 "src/sonyflake_turbo/sonyflake.c",
                 "src/sonyflake_turbo/machine_ids.c",
                 "src/sonyflake_turbo/sleep_wrapper.c",
+                "src/sonyflake_turbo/async.c",
             ],
             define_macros=define_macros,
             py_limited_api=py_limited_api,

src/sonyflake_turbo/__init__.py

@@ -6,6 +6,7 @@
     SONYFLAKE_SEQUENCE_BITS,
     SONYFLAKE_SEQUENCE_MAX,
     SONYFLAKE_TIME_OFFSET,
+    AsyncSonyFlake,
     MachineIDLCG,
     SonyFlake,
 )
@@ -18,6 +19,7 @@
     "SONYFLAKE_SEQUENCE_BITS",
     "SONYFLAKE_SEQUENCE_MAX",
     "SONYFLAKE_TIME_OFFSET",
+    "AsyncSonyFlake",
     "MachineIDLCG",
     "SonyFlake",
 ]

src/sonyflake_turbo/_sonyflake.c

@@ -5,12 +5,14 @@
 
 #include "module.h"
 #include "sonyflake.h"
+#include "async.h"
 #include "machine_ids.h"
 #include "sleep_wrapper.h"
 
 static int sonyflake_module_traverse(PyObject *m, visitproc visit, void *arg) {
 	struct sonyflake_module_state *state = PyModule_GetState(m);
 	Py_VISIT(state->sonyflake_cls);
+	Py_VISIT(state->async_sonyflake_cls);
 	Py_VISIT(state->machine_id_lcg_cls);
 	Py_VISIT(state->sleep_wrapper_cls);
 	return 0;
@@ -19,6 +21,7 @@ static int sonyflake_module_traverse(PyObject *m, visitproc visit, void *arg) {
 static int sonyflake_module_clear(PyObject *m) {
 	struct sonyflake_module_state *state = PyModule_GetState(m);
 	Py_CLEAR(state->sonyflake_cls);
+	Py_CLEAR(state->async_sonyflake_cls);
 	Py_CLEAR(state->machine_id_lcg_cls);
 	Py_CLEAR(state->sleep_wrapper_cls);
 	return 0;
@@ -62,7 +65,9 @@ static int sonyflake_exec(PyObject *module) {
 	struct sonyflake_module_state *state = PyModule_GetState(module);
 
 	state->sonyflake_cls = NULL;
+	state->async_sonyflake_cls = NULL;
 	state->machine_id_lcg_cls = NULL;
+	state->sleep_wrapper_cls = NULL;
 
 	state->sonyflake_cls = PyType_FromModuleAndSpec(module, &sonyflake_type_spec, NULL);
 
@@ -74,6 +79,16 @@ static int sonyflake_exec(PyObject *module) {
 		goto err;
 	}
 
+	state->async_sonyflake_cls = PyType_FromModuleAndSpec(module, &async_sonyflake_type_spec, NULL);
+
+	if (!state->async_sonyflake_cls) {
+		goto err;
+	}
+
+	if (PyModule_AddObjectRef(module, "AsyncSonyFlake", state->async_sonyflake_cls) < 0) {
+		goto err;
+	}
+
 	state->machine_id_lcg_cls = PyType_FromModuleAndSpec(module, &machine_id_lcg_spec, NULL);
 
 	if (!state->machine_id_lcg_cls) {
@@ -107,6 +122,7 @@ static int sonyflake_exec(PyObject *module) {
 err:
 	Py_CLEAR(state->sleep_wrapper_cls);
 	Py_CLEAR(state->machine_id_lcg_cls);
+	Py_CLEAR(state->async_sonyflake_cls);
 	Py_CLEAR(state->sonyflake_cls);
 
 	return -1;

src/sonyflake_turbo/_sonyflake.pyi

@@ -1,4 +1,4 @@
-from typing import List, Optional, Callable, Awaitable, TypeAlias, TypeVar
+from typing import Awaitable, Callable, Generator, List, Optional, TypeAlias, TypeVar
 
 try:
     from typing import Self
@@ -15,7 +15,6 @@ SONYFLAKE_TIME_OFFSET: int
 AsyncSleep: TypeAlias = Callable[[float], Awaitable[None]]
 T = TypeVar("T")
 
-
 async def sleep_wrapper(obj: T, sleep: AsyncSleep, to_sleep: float) -> T:
     """C version of:
 
@@ -56,17 +55,97 @@ class SonyFlake:
         time.
 
         Important:
-            The more ids you request, the more other threads has to wait
-            upon next :meth:`next` or :meth:`n` call.
+            The more ids you request, the more other threads have to wait
+            upon the next id(s).
+
+        Args:
+            n: Number of ids to generate. Must be greater than 0.
+
+        Raises:
+            ValueError: if n <= 0
+
+        Returns:
+            List of ids.
+        """
+
+class AsyncSonyFlake:
+    """Async wrapper for :class:`SonyFlake`.
+
+    Implements Awaitable and AsyncIterator protocols.
+
+    Important:
+        Main state is stored in SonyFlake. This class has minimum logic on its
+        own, the only difference is that instead of doing thread-blocking sleep,
+        it is delegated to the provided ``sleep``. Instance of
+        class:`sleep_wrapper` is returned in ``__call__``, ``__anext__`` and
+        ``__await``.
+
+    Usage:
+
+        .. code-block:: python
+
+            import asyncio
+
+            sf = SonyFlake(0x1337, 0xCAFE, start_time=1749081600)
+            asf = AsyncSonyFlake(sf, asyncio.sleep)
+
+            print(await asf)
+            print(await asf(5))
+
+            async for id_ in asf:
+                print(id_)
+                break  # AsyncSonyFlake is an infinite generator
+    """
+
+    def __init__(self, sf: SonyFlake, sleep: AsyncSleep) -> None:
+        """Initialize AsyncSonyFlake ID generator.
+
+        Args:
+            sf: Instance of the :class:`SonyFlake`.
+            sleep: Either `asyncio.sleep` or `trio.sleep`.
+
+        Raises:
+            ValueError: Invalid values of ``machine_id`` or ``start_time``.
+            TypeError: ``machine_id`` or ``start_time`` are not integers.
+        """
+
+    def __call__(self, n: int) -> Awaitable[list[int]]:
+        """Generate multiple SonyFlake IDs at once.
+
+        Roughly equivalent to `[await asf for _ in range(n)]`, but more
+        efficient. This method saves on task/context switches and syscalls for
+        getting current time.
+
+        Important:
+            The more ids you request, the more other coroutines have to wait
+            upon the next id(s).
 
         Args:
             n: Number of ids to generate. Must be greater than 0.
 
+        Raises:
+            ValueError: if n <= 0
+
         Returns:
             List of ids.
         """
 
+    def __await__(self) -> Generator[None, None, int]:
+        """Produce a SonyFlake ID."""
+
+    def __anext__(self) -> Awaitable[int]:
+        """Produce a SonyFlake ID."""
+
+    def __aiter__(self) -> Self:
+        """Returns ``self``."""
+
 class MachineIDLCG:
+    """A simple LCG producing ints suitable to be used as ``machine_id``.
+
+    Intended to be used in examples, tests, or when concurrency is not an issue.
+    You generally want only one instance of this class per process.
+    """
+
     def __new__(cls, seed: int, /) -> Self:
         """Make a LCG.
 

src/sonyflake_turbo/async.c

@@ -0,0 +1,148 @@
+#include <Python.h>
+
+#include "async.h"
+#include "sleep_wrapper.h"
+#include "sonyflake.h"
+
+
+struct async_sonyflake_state {
+	PyObject_HEAD
+
+	struct sonyflake_state *sf;
+	PyObject *sleep;
+};
+
+static int async_sonyflake_clear(PyObject *py_self) {
+	struct async_sonyflake_state *self = (struct async_sonyflake_state *) py_self;
+	Py_CLEAR(self->sf);
+	Py_CLEAR(self->sleep);
+	return 0;
+}
+
+static void async_sonyflake_dealloc(PyObject *py_self) {
+	async_sonyflake_clear(py_self);
+
+	PyTypeObject *tp = Py_TYPE(py_self);
+	freefunc tp_free = PyType_GetSlot(tp, Py_tp_free);
+
+	assert(tp_free != NULL);
+
+	tp_free(py_self);
+	Py_DECREF(tp);
+}
+
+static int async_sonyflake_traverse(PyObject *py_self, visitproc visit, void *arg) {
+	struct async_sonyflake_state *self = (struct async_sonyflake_state *) py_self;
+	Py_VISIT(self->sf);
+	Py_VISIT(self->sleep);
+	Py_VISIT(Py_TYPE(py_self));
+	return 0;
+}
+
+static PyObject *async_sonyflake_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject *Py_UNUSED(kwargs)) {
+	allocfunc tp_alloc = PyType_GetSlot(type, Py_tp_alloc);
+
+	assert(tp_alloc != NULL);
+
+	struct async_sonyflake_state *self = (void *) tp_alloc(type, 0);
+
+	if (!self) {
+		return NULL;
+	}
+
+	self->sf = NULL;
+	self->sleep = NULL;
+
+	return (PyObject *) self;
+}
+
+static int async_sonyflake_init(PyObject *py_self, PyObject *args, PyObject *Py_UNUSED(kwargs)) {
+	struct async_sonyflake_state *self = (struct async_sonyflake_state *) py_self;
+
+	if (!PyArg_ParseTuple(args, "OO", &self->sf, &self->sleep)) {
+		return -1;
+	}
+
+	struct sonyflake_module_state *module = PyType_GetModuleState(Py_TYPE(self));
+
+	assert(module != NULL);
+
+	if (!PyObject_IsInstance((PyObject *) self->sf, module->sonyflake_cls)) {
+		PyErr_SetString(PyExc_TypeError, "sf must be instance of SonyFlake");
+		return -1;
+	}
+
+	if (!PyCallable_Check(self->sleep)) {
+		PyErr_SetString(PyExc_TypeError, "sleep must be callable");
+		return -1;
+	}
+
+	Py_INCREF(self->sf);
+	Py_INCREF(self->sleep);
+
+	return 0;
+}
+
+static inline PyObject *wrap_into_sleep(
+	struct async_sonyflake_state *self,
+	struct sonyflake_next_sleep_info *sleep_info,
+	PyObject *obj
+) {
+	if (!obj) {
+		return NULL;
+	}
+
+	struct sonyflake_module_state *module = PyType_GetModuleState(Py_TYPE(self));
+
+	assert(module);
+	assert(sleep_info);
+
+	struct timespec diff = sleep_info->future;
+
+	sub_diff(&diff, &sleep_info->now);
+
+	PyObject *coro = PyObject_CallFunction(
+		module->sleep_wrapper_cls,
+		"OOd",
+		obj,
+		self->sleep,
+		timespec_to_double(&diff)
+	);
+
+	Py_DECREF(obj);
+
+	return coro;
+}
+
+static PyObject *async_sonyflake_call(struct async_sonyflake_state *self, PyObject *args) {
+	struct sonyflake_next_sleep_info sleep_info;
+
+	return wrap_into_sleep(self, &sleep_info, sonyflake_next_py(self->sf, args, &sleep_info));
+}
+
+static PyObject *async_sonyflake_await(struct async_sonyflake_state *self) {
+	struct sonyflake_next_sleep_info sleep_info;
+
+	return wrap_into_sleep(self, &sleep_info, sonyflake_next_py(self->sf, NULL, &sleep_info));
+}
+
+PyType_Slot async_sonyflake_type_slots[] = {
+	{Py_tp_alloc, PyType_GenericAlloc},
+	{Py_tp_dealloc, async_sonyflake_dealloc},
+	{Py_tp_traverse, async_sonyflake_traverse},
+	{Py_tp_clear, async_sonyflake_clear},
+	{Py_tp_new, async_sonyflake_new},
+	{Py_tp_init, async_sonyflake_init},
+	{Py_tp_call, async_sonyflake_call},
+	{Py_am_await, async_sonyflake_await},
+	{Py_am_anext, async_sonyflake_await},
+	{Py_am_aiter, PyObject_SelfIter},
+	{0, 0},
+};
+
+PyType_Spec async_sonyflake_type_spec = {
+	.name = MODULE_NAME ".AsyncSonyFlake",
+	.basicsize = sizeof(struct async_sonyflake_state),
+	.flags = Py_TPFLAGS_DEFAULT,
+	.slots = async_sonyflake_type_slots,
+};

src/sonyflake_turbo/async.h

@@ -0,0 +1,7 @@
+#pragma once
+
+#include "sonyflake.h"
+
+struct async_sonyflake_state;
+
+extern PyType_Spec async_sonyflake_type_spec;

src/sonyflake_turbo/module.h

@@ -8,6 +8,7 @@
 
 struct sonyflake_module_state {
 	PyObject *sonyflake_cls;
+	PyObject *async_sonyflake_cls;
 	PyObject *machine_id_lcg_cls;
 	PyObject *sleep_wrapper_cls;
 };