Skip to content

Commit 4576ba1

Browse files
ZipFileZipFile
committed
Simplify async implementation
1 parent e8a1555 commit 4576ba1

14 files changed

Lines changed: 184 additions & 542 deletions

‎setup.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,6 @@
3333
"src/sonyflake_turbo/_sonyflake.c",
3434
"src/sonyflake_turbo/sonyflake.c",
3535
"src/sonyflake_turbo/machine_ids.c",
36-
"src/sonyflake_turbo/sleep_wrapper.c",
37-
"src/sonyflake_turbo/async.c",
3836
],
3937
define_macros=define_macros,
4038
py_limited_api=py_limited_api,

‎src/sonyflake_turbo/__init__.py‎

Lines changed: 95 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,14 @@
1+
from typing import (
2+
Any,
3+
AsyncIterable,
4+
AsyncIterator,
5+
Awaitable,
6+
Callable,
7+
ClassVar,
8+
Generator,
9+
TypeAlias,
10+
)
11+
112
from ._sonyflake import (
213
SONYFLAKE_EPOCH,
314
SONYFLAKE_MACHINE_ID_BITS,
@@ -6,11 +17,12 @@
617
SONYFLAKE_SEQUENCE_BITS,
718
SONYFLAKE_SEQUENCE_MAX,
819
SONYFLAKE_TIME_OFFSET,
9-
AsyncSonyFlake,
1020
MachineIDLCG,
1121
SonyFlake,
1222
)
1323

24+
AsyncSleep: TypeAlias = Callable[[float], Awaitable[None]]
25+
1426
__all__ = [
1527
"SONYFLAKE_EPOCH",
1628
"SONYFLAKE_MACHINE_ID_BITS",
@@ -23,3 +35,85 @@
2335
"MachineIDLCG",
2436
"SonyFlake",
2537
]
38+
39+
40+
class AsyncSonyFlake(Awaitable[int], AsyncIterable[int]):
41+
"""Async wrapper for :class:`SonyFlake`.
42+
43+
Usage:
44+
45+
.. code-block:: python
46+
47+
import asyncio
48+
49+
sf = SonyFlake(0x1337, 0xCAFE, start_time=1749081600)
50+
asf = AsyncSonyFlake(sf, asyncio.sleep)
51+
52+
print(await asf)
53+
print(await asf(5))
54+
55+
async for id_ in asf:
56+
print(id_)
57+
break
58+
"""
59+
60+
__slots__: ClassVar[tuple[str, ...]] = ("sf", "sleep")
61+
sleep: AsyncSleep
62+
63+
def __init__(self, sf: SonyFlake, sleep: AsyncSleep | None = None) -> None:
64+
"""Initialize AsyncSonyFlake ID generator.
65+
66+
Args:
67+
sf: Instance of the :class:`SonyFlake`.
68+
sleep: Either `asyncio.sleep` or `trio.sleep`.
69+
"""
70+
71+
if sleep is None:
72+
from asyncio import sleep
73+
74+
assert sleep is not None
75+
76+
self.sf = sf
77+
self.sleep = sleep
78+
79+
async def __call__(self, n: int) -> list[int]:
80+
"""Generate multiple SonyFlake IDs at once.
81+
82+
Roughly equivalent to ``[await asf for _ in range(n)]``, but more
83+
efficient. This method saves on task/context switches and syscalls for
84+
getting current time.
85+
86+
Args:
87+
n: Number of ids to generate.
88+
89+
Returns:
90+
List of ids.
91+
"""
92+
93+
ids, to_sleep = self.sf._raw(n)
94+
95+
await self.sleep(to_sleep)
96+
97+
return ids
98+
99+
def __await__(self) -> Generator[Any, Any, int]:
100+
"""Produce a SonyFlake ID."""
101+
102+
id_, to_sleep = self.sf._raw(None)
103+
104+
yield from self.sleep(to_sleep).__await__()
105+
106+
return id_
107+
108+
def __aiter__(self) -> AsyncIterator[int]:
109+
"""Return an infinite SonyFlake ID async iterator."""
110+
111+
return self._gen().__aiter__()
112+
113+
async def _gen(self) -> AsyncIterator[int]:
114+
"""Infinite SonyFlake ID async generator."""
115+
116+
while True:
117+
id_, to_sleep = self.sf._raw(None)
118+
await self.sleep(to_sleep)
119+
yield id_

‎src/sonyflake_turbo/_sonyflake.c‎

Lines changed: 0 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -5,25 +5,19 @@
55

66
#include "module.h"
77
#include "sonyflake.h"
8-
#include "async.h"
98
#include "machine_ids.h"
10-
#include "sleep_wrapper.h"
119

1210
static int sonyflake_module_traverse(PyObject *m, visitproc visit, void *arg) {
1311
struct sonyflake_module_state *state = PyModule_GetState(m);
1412
Py_VISIT(state->sonyflake_cls);
15-
Py_VISIT(state->async_sonyflake_cls);
1613
Py_VISIT(state->machine_id_lcg_cls);
17-
Py_VISIT(state->sleep_wrapper_cls);
1814
return 0;
1915
}
2016

2117
static int sonyflake_module_clear(PyObject *m) {
2218
struct sonyflake_module_state *state = PyModule_GetState(m);
2319
Py_CLEAR(state->sonyflake_cls);
24-
Py_CLEAR(state->async_sonyflake_cls);
2520
Py_CLEAR(state->machine_id_lcg_cls);
26-
Py_CLEAR(state->sleep_wrapper_cls);
2721
return 0;
2822
}
2923

@@ -65,9 +59,7 @@ static int sonyflake_exec(PyObject *module) {
6559
struct sonyflake_module_state *state = PyModule_GetState(module);
6660

6761
state->sonyflake_cls = NULL;
68-
state->async_sonyflake_cls = NULL;
6962
state->machine_id_lcg_cls = NULL;
70-
state->sleep_wrapper_cls = NULL;
7163

7264
state->sonyflake_cls = PyType_FromModuleAndSpec(module, &sonyflake_type_spec, NULL);
7365

@@ -79,16 +71,6 @@ static int sonyflake_exec(PyObject *module) {
7971
goto err;
8072
}
8173

82-
state->async_sonyflake_cls = PyType_FromModuleAndSpec(module, &async_sonyflake_type_spec, NULL);
83-
84-
if (!state->async_sonyflake_cls) {
85-
goto err;
86-
}
87-
88-
if (PyModule_AddObjectRef(module, "AsyncSonyFlake", state->async_sonyflake_cls) < 0) {
89-
goto err;
90-
}
91-
9274
state->machine_id_lcg_cls = PyType_FromModuleAndSpec(module, &machine_id_lcg_spec, NULL);
9375

9476
if (!state->machine_id_lcg_cls) {
@@ -99,16 +81,6 @@ static int sonyflake_exec(PyObject *module) {
9981
goto err;
10082
}
10183

102-
state->sleep_wrapper_cls = PyType_FromModuleAndSpec(module, &sleep_wrapper_type_spec, NULL);
103-
104-
if (!state->sleep_wrapper_cls) {
105-
goto err;
106-
}
107-
108-
if (PyModule_AddObjectRef(module, "sleep_wrapper", state->sleep_wrapper_cls) < 0) {
109-
goto err;
110-
}
111-
11284
PyModule_AddIntMacro(module, SONYFLAKE_EPOCH);
11385
PyModule_AddIntMacro(module, SONYFLAKE_SEQUENCE_BITS);
11486
PyModule_AddIntMacro(module, SONYFLAKE_SEQUENCE_MAX);
@@ -120,9 +92,7 @@ static int sonyflake_exec(PyObject *module) {
12092
return 0;
12193

12294
err:
123-
Py_CLEAR(state->sleep_wrapper_cls);
12495
Py_CLEAR(state->machine_id_lcg_cls);
125-
Py_CLEAR(state->async_sonyflake_cls);
12696
Py_CLEAR(state->sonyflake_cls);
12797

12898
return -1;

‎src/sonyflake_turbo/_sonyflake.pyi‎

Lines changed: 7 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Awaitable, Callable, Generator, List, Optional, TypeAlias, TypeVar
1+
from typing import overload
22

33
try:
44
from typing import Self
@@ -12,19 +12,9 @@ SONYFLAKE_MACHINE_ID_BITS: int
1212
SONYFLAKE_MACHINE_ID_MAX: int
1313
SONYFLAKE_MACHINE_ID_OFFSET: int
1414
SONYFLAKE_TIME_OFFSET: int
15-
AsyncSleep: TypeAlias = Callable[[float], Awaitable[None]]
16-
T = TypeVar("T")
17-
18-
async def sleep_wrapper(obj: T, sleep: AsyncSleep, to_sleep: float) -> T:
19-
"""C version of:
20-
21-
async def sleep_wrapper(obj, sleep, to_sleep):
22-
await sleep(to_sleep)
23-
return obj
24-
"""
2515

2616
class SonyFlake:
27-
def __init__(self, *machine_id: int, start_time: Optional[int] = None):
17+
def __init__(self, *machine_id: int, start_time: int | None = None):
2818
"""Initialize SonyFlake ID generator.
2919
3020
Args:
@@ -47,7 +37,7 @@ class SonyFlake:
4737
def __next__(self) -> int:
4838
"""Produce a SonyFlake ID."""
4939

50-
def __call__(self, n: int, /) -> List[int]:
40+
def __call__(self, n: int, /) -> list[int]:
5141
"""Generate multiple SonyFlake IDs at once.
5242
5343
Roughly equivalent to `[next(sf) for _ in range(n)]`, but more
@@ -68,76 +58,10 @@ class SonyFlake:
6858
List of ids.
6959
"""
7060

71-
class AsyncSonyFlake:
72-
"""Async wrapper for :class:`SonyFlake`.
73-
74-
Implements Awaitable and AsyncIterator protocols.
75-
76-
Important:
77-
Main state is stored in SonyFlake. This class has minimum logic on its
78-
own, the only difference is that instead of doing thread-blocking sleep,
79-
it is delegated to the provided ``sleep``. Instance of
80-
class:`sleep_wrapper` is returned in ``__call__``, ``__anext__`` and
81-
``__await``.
82-
83-
Usage:
84-
85-
.. code-block:: python
86-
87-
import asyncio
88-
89-
sf = SonyFlake(0x1337, 0xCAFE, start_time=1749081600)
90-
asf = AsyncSonyFlake(sf, asyncio.sleep)
91-
92-
print(await asf)
93-
print(await asf(5))
94-
95-
async for id_ in asf:
96-
print(id_)
97-
break # AsyncSonyFlake is an infinite generator
98-
"""
99-
100-
def __init__(self, sf: SonyFlake, sleep: AsyncSleep) -> None:
101-
"""Initialize AsyncSonyFlake ID generator.
102-
103-
Args:
104-
sf: Instance of the :class:`SonyFlake`.
105-
sleep: Either `asyncio.sleep` or `trio.sleep`.
106-
107-
Raises:
108-
ValueError: Invalid values of ``machine_id`` or ``start_time``.
109-
TypeError: ``machine_id`` or ``start_time`` are not integers.
110-
"""
111-
112-
def __call__(self, n: int) -> Awaitable[list[int]]:
113-
"""Generate multiple SonyFlake IDs at once.
114-
115-
Roughly equivalent to `[await asf for _ in range(n)]`, but more
116-
efficient. This method saves on task/context switches and syscalls for
117-
getting current time.
118-
119-
Important:
120-
The more ids you request, the more other coroutines have to wait
121-
upon the next id(s).
122-
123-
Args:
124-
n: Number of ids to generate. Must be greater than 0.
125-
126-
Raises:
127-
ValueError: if n <= 0
128-
129-
Returns:
130-
List of ids.
131-
"""
132-
133-
def __await__(self) -> Generator[None, None, int]:
134-
"""Produce a SonyFlake ID."""
135-
136-
def __anext__(self) -> Awaitable[int]:
137-
"""Produce a SonyFlake ID."""
138-
139-
def __aiter__(self) -> Self:
140-
"""Returns ``self``."""
61+
@overload
62+
def _raw(self, n: int) -> tuple[list[int], float]: ...
63+
@overload
64+
def _raw(self, n: None) -> tuple[int, float]: ...
14165

14266
class MachineIDLCG:
14367
"""A simple LCG producing ints suitable to be used as ``machine_id``.

0 commit comments

Comments
 (0)