Skip to content

Commit d7a794f

Browse files
ZipFileZipFile
committed
Implement sleep wrapper
1 parent 91c2007 commit d7a794f

6 files changed

Lines changed: 257 additions & 1 deletion

File tree

‎setup.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
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",
3637
],
3738
define_macros=define_macros,
3839
py_limited_api=py_limited_api,

‎src/sonyflake_turbo/_sonyflake.c‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,21 @@
66
#include "module.h"
77
#include "sonyflake.h"
88
#include "machine_ids.h"
9+
#include "sleep_wrapper.h"
910

1011
static int sonyflake_module_traverse(PyObject *m, visitproc visit, void *arg) {
1112
struct sonyflake_module_state *state = PyModule_GetState(m);
1213
Py_VISIT(state->sonyflake_cls);
1314
Py_VISIT(state->machine_id_lcg_cls);
15+
Py_VISIT(state->sleep_wrapper_cls);
1416
return 0;
1517
}
1618

1719
static int sonyflake_module_clear(PyObject *m) {
1820
struct sonyflake_module_state *state = PyModule_GetState(m);
1921
Py_CLEAR(state->sonyflake_cls);
2022
Py_CLEAR(state->machine_id_lcg_cls);
23+
Py_CLEAR(state->sleep_wrapper_cls);
2124
return 0;
2225
}
2326

@@ -81,6 +84,16 @@ static int sonyflake_exec(PyObject *module) {
8184
goto err;
8285
}
8386

87+
state->sleep_wrapper_cls = PyType_FromModuleAndSpec(module, &sleep_wrapper_type_spec, NULL);
88+
89+
if (!state->sleep_wrapper_cls) {
90+
goto err;
91+
}
92+
93+
if (PyModule_AddObjectRef(module, "sleep_wrapper", state->sleep_wrapper_cls) < 0) {
94+
goto err;
95+
}
96+
8497
PyModule_AddIntMacro(module, SONYFLAKE_EPOCH);
8598
PyModule_AddIntMacro(module, SONYFLAKE_SEQUENCE_BITS);
8699
PyModule_AddIntMacro(module, SONYFLAKE_SEQUENCE_MAX);
@@ -92,6 +105,7 @@ static int sonyflake_exec(PyObject *module) {
92105
return 0;
93106

94107
err:
108+
Py_CLEAR(state->sleep_wrapper_cls);
95109
Py_CLEAR(state->machine_id_lcg_cls);
96110
Py_CLEAR(state->sonyflake_cls);
97111

‎src/sonyflake_turbo/_sonyflake.pyi‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import List, Optional
1+
from typing import List, Optional, Callable, Awaitable, TypeAlias, TypeVar
22

33
try:
44
from typing import Self
@@ -12,6 +12,17 @@ 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+
19+
async def sleep_wrapper(obj: T, sleep: AsyncSleep, to_sleep: float) -> T:
20+
"""C version of:
21+
22+
async def sleep_wrapper(obj, sleep, to_sleep):
23+
await sleep(to_sleep)
24+
return obj
25+
"""
1526

1627
class SonyFlake:
1728
def __init__(self, *machine_id: int, start_time: Optional[int] = None):

‎src/sonyflake_turbo/sleep_wrapper.c‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
#include <Python.h>
2+
3+
#include "module.h"
4+
#include "sleep_wrapper.h"
5+
6+
struct sleep_wrapper_state {
7+
PyObject_HEAD
8+
9+
PyObject *obj;
10+
PyObject *sleep;
11+
PyObject *to_sleep;
12+
PyObject *it;
13+
};
14+
15+
static int sleep_wrapper_clear(PyObject *py_self) {
16+
struct sleep_wrapper_state *self = (struct sleep_wrapper_state *) py_self;
17+
Py_CLEAR(self->obj);
18+
Py_CLEAR(self->sleep);
19+
Py_CLEAR(self->to_sleep);
20+
Py_CLEAR(self->it);
21+
return 0;
22+
}
23+
24+
static void sleep_wrapper_dealloc(PyObject *py_self) {
25+
sleep_wrapper_clear(py_self);
26+
27+
PyTypeObject *tp = Py_TYPE(py_self);
28+
freefunc tp_free = PyType_GetSlot(tp, Py_tp_free);
29+
30+
assert(tp_free != NULL);
31+
32+
tp_free(py_self);
33+
Py_DECREF(tp);
34+
}
35+
36+
static int sleep_wrapper_traverse(PyObject *py_self, visitproc visit, void *arg) {
37+
struct sleep_wrapper_state *self = (struct sleep_wrapper_state *) py_self;
38+
Py_VISIT(self->obj);
39+
Py_VISIT(self->sleep);
40+
Py_VISIT(self->to_sleep);
41+
Py_VISIT(self->it);
42+
Py_VISIT(Py_TYPE(py_self));
43+
return 0;
44+
}
45+
46+
static PyObject *sleep_wrapper_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject *Py_UNUSED(kwargs)) {
47+
allocfunc tp_alloc = PyType_GetSlot(type, Py_tp_alloc);
48+
49+
assert(tp_alloc != NULL);
50+
51+
struct sleep_wrapper_state *self = (void *) tp_alloc(type, 0);
52+
53+
if (!self) {
54+
return NULL;
55+
}
56+
57+
self->obj = NULL;
58+
self->sleep = NULL;
59+
self->to_sleep = NULL;
60+
self->it = NULL;
61+
62+
return (PyObject *) self;
63+
}
64+
65+
static int sleep_wrapper_init(PyObject *py_self, PyObject *args, PyObject *Py_UNUSED(kwargs)) {
66+
struct sleep_wrapper_state *self = (struct sleep_wrapper_state *) py_self;
67+
68+
if (!PyArg_ParseTuple(args, "OOO", &self->obj, &self->sleep, &self->to_sleep)) {
69+
return -1;
70+
}
71+
72+
Py_INCREF(self->obj);
73+
Py_INCREF(self->sleep);
74+
Py_INCREF(self->to_sleep);
75+
76+
return 0;
77+
}
78+
79+
static PyObject *sleep_wrapper_iternext(struct sleep_wrapper_state *self) {
80+
if (self->it) {
81+
PyObject *obj = PyIter_Next(self->it);
82+
83+
if (PyErr_Occurred()) {
84+
return NULL;
85+
} else if (obj) {
86+
return obj;
87+
}
88+
89+
Py_CLEAR(self->it);
90+
}
91+
92+
PyErr_SetObject(PyExc_StopIteration, self->obj);
93+
94+
return NULL;
95+
}
96+
97+
static PySendResult sleep_wrapper_send(PyObject *py_self, PyObject *arg, PyObject **presult) {
98+
struct sleep_wrapper_state *self = (struct sleep_wrapper_state *) py_self;
99+
100+
if (self->it) {
101+
PySendResult result = PyIter_Send(self->it, arg, presult);
102+
103+
if (result != PYGEN_RETURN) {
104+
return result;
105+
}
106+
107+
Py_CLEAR(self->it);
108+
}
109+
110+
PyErr_SetObject(PyExc_StopIteration, self->obj);
111+
112+
return PYGEN_ERROR;
113+
}
114+
115+
static PyObject *sleep_wrapper_send_meth(PyObject *self, PyObject *arg)
116+
{
117+
PyObject *result = NULL;
118+
119+
if (sleep_wrapper_send(self, arg, &result) == PYGEN_RETURN) {
120+
return result;
121+
}
122+
123+
return NULL;
124+
}
125+
126+
static PyObject *sleep_wrapper_iter_or_await(PyObject *py_self) {
127+
struct sleep_wrapper_state *self = (struct sleep_wrapper_state *) py_self;
128+
129+
Py_CLEAR(self->it);
130+
131+
PyObject *coro = PyObject_CallFunction(self->sleep, "O", self->to_sleep);
132+
133+
if (!coro) {
134+
return NULL;
135+
}
136+
137+
unaryfunc am_await = PyType_GetSlot(Py_TYPE(coro), Py_am_await);
138+
139+
if (!am_await) {
140+
PyErr_SetString(PyExc_TypeError, "sleep() does not support __await__");
141+
Py_DECREF(coro);
142+
return NULL;
143+
}
144+
145+
self->it = am_await(coro);
146+
147+
Py_DECREF(coro);
148+
149+
if (!self->it) {
150+
return NULL;
151+
}
152+
153+
if (!PyIter_Check(self->it)) {
154+
PyErr_SetString(PyExc_TypeError, "sleep().__await__() did not return an iterator");
155+
Py_CLEAR(self->it);
156+
return NULL;
157+
}
158+
159+
Py_INCREF(py_self);
160+
161+
return py_self;
162+
}
163+
164+
static PyMethodDef sleep_wrapper_methods[] = {
165+
{"send", sleep_wrapper_send_meth, METH_O, ""},
166+
{NULL, NULL, 0, NULL}
167+
};
168+
169+
PyType_Slot sleep_wrapper_type_slots[] = {
170+
{Py_tp_alloc, PyType_GenericAlloc},
171+
{Py_tp_dealloc, sleep_wrapper_dealloc},
172+
{Py_tp_traverse, sleep_wrapper_traverse},
173+
{Py_tp_clear, sleep_wrapper_clear},
174+
{Py_tp_iter, sleep_wrapper_iter_or_await},
175+
{Py_tp_iternext, sleep_wrapper_iternext},
176+
{Py_tp_new, sleep_wrapper_new},
177+
{Py_tp_init, sleep_wrapper_init},
178+
{Py_am_await, sleep_wrapper_iter_or_await},
179+
{Py_am_send, sleep_wrapper_send},
180+
{Py_tp_methods, sleep_wrapper_methods},
181+
{0, 0},
182+
};
183+
184+
PyType_Spec sleep_wrapper_type_spec = {
185+
.name = MODULE_NAME ".sleep_wrapper",
186+
.basicsize = sizeof(struct sleep_wrapper_state),
187+
.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
188+
.slots = sleep_wrapper_type_slots,
189+
};

‎src/sonyflake_turbo/sleep_wrapper.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
#pragma once
2+
3+
#include <Python.h>
4+
5+
struct sleep_wrapper_state;
6+
extern PyType_Spec sleep_wrapper_type_spec;

‎tests/test_sleep_wrapper.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
from typing import Awaitable, Callable
2+
3+
from pytest import mark
4+
5+
from sonyflake_turbo._sonyflake import sleep_wrapper
6+
7+
AsyncSleep = Callable[[float], Awaitable[None]]
8+
obj = object()
9+
to_sleep = 0.1
10+
11+
12+
async def _test_sleep_wrapper(sleep: AsyncSleep) -> None:
13+
sleep_called_with = None
14+
15+
def _sleep(t: float) -> Awaitable[None]:
16+
nonlocal sleep_called_with
17+
sleep_called_with = t
18+
return sleep(t)
19+
20+
assert await sleep_wrapper(obj, _sleep, to_sleep) is obj
21+
assert sleep_called_with is to_sleep, "sleep not called"
22+
23+
24+
@mark.asyncio
25+
async def test_asyncio_sleep() -> None:
26+
from asyncio import sleep
27+
28+
await _test_sleep_wrapper(sleep)
29+
30+
31+
@mark.trio
32+
async def test_trio_gen() -> None:
33+
from trio import sleep
34+
35+
await _test_sleep_wrapper(sleep)

0 commit comments

Comments
 (0)