Add pure python fallback implementation

ZipFile · ZipFile · commit cebd25b9c370 · 2026-03-04T12:28:01.000Z
diff --git a/setup.py b/setup.py
@@ -1,14 +1,20 @@
 #!/usr/bin/env python
+import os
+import sys
 import sysconfig
 from typing import Any, Dict, List, Optional, Tuple
 
 from setuptools import Extension, setup
 
+build: bool = os.environ.get("SONYFLAKE_TURBO_BUILD", "1").lower() in ("1", "true")
 options: Dict[str, Any] = {}
 define_macros: List[Tuple[str, Optional[str]]] = []
 py_limited_api: bool = not sysconfig.get_config_var("Py_GIL_DISABLED")
 cflags: List[str] = []
 
+if sys.implementation.name != "cpython":
+    build = False
+
 if sysconfig.get_platform().startswith("win"):
     cflags.append("/utf-8")
     cflags.append("/std:c17")
@@ -24,9 +30,10 @@
     options["bdist_wheel"] = {"py_limited_api": "cp310"}
     define_macros.append(("Py_LIMITED_API", "0x030a00f0"))
 
-setup(
-    options=options,
-    ext_modules=[
+
+setup_kwargs = {
+    "options": options,
+    "ext_modules": [
         Extension(
             "sonyflake_turbo._sonyflake",
             sources=[
@@ -39,4 +46,9 @@
             extra_compile_args=cflags,
         ),
     ],
-)
+}
+
+if not build:
+    setup_kwargs = {}
+
+setup(**setup_kwargs)
diff --git a/src/sonyflake_turbo/__init__.py b/src/sonyflake_turbo/__init__.py
@@ -1,27 +1,27 @@
-from typing import (
-    Any,
-    AsyncIterable,
-    AsyncIterator,
-    Awaitable,
-    Callable,
-    ClassVar,
-    Generator,
-    TypeAlias,
-)
-
-from ._sonyflake import (
-    SONYFLAKE_EPOCH,
-    SONYFLAKE_MACHINE_ID_BITS,
-    SONYFLAKE_MACHINE_ID_MAX,
-    SONYFLAKE_MACHINE_ID_OFFSET,
-    SONYFLAKE_SEQUENCE_BITS,
-    SONYFLAKE_SEQUENCE_MAX,
-    SONYFLAKE_TIME_OFFSET,
-    MachineIDLCG,
-    SonyFlake,
-)
-
-AsyncSleep: TypeAlias = Callable[[float], Awaitable[None]]
+try:
+    from ._sonyflake import (
+        SONYFLAKE_EPOCH,
+        SONYFLAKE_MACHINE_ID_BITS,
+        SONYFLAKE_MACHINE_ID_MAX,
+        SONYFLAKE_MACHINE_ID_OFFSET,
+        SONYFLAKE_SEQUENCE_BITS,
+        SONYFLAKE_SEQUENCE_MAX,
+        SONYFLAKE_TIME_OFFSET,
+        MachineIDLCG,
+        SonyFlake,
+    )
+except ImportError:  # pragma: no cover
+    from .pure import (
+        SONYFLAKE_EPOCH,
+        SONYFLAKE_MACHINE_ID_BITS,
+        SONYFLAKE_MACHINE_ID_MAX,
+        SONYFLAKE_MACHINE_ID_OFFSET,
+        SONYFLAKE_SEQUENCE_BITS,
+        SONYFLAKE_SEQUENCE_MAX,
+        SONYFLAKE_TIME_OFFSET,
+        MachineIDLCG,
+        SonyFlake,
+    )
 
 __all__ = [
     "SONYFLAKE_EPOCH",
@@ -37,7 +37,7 @@
 ]
 
 
-class AsyncSonyFlake(Awaitable[int], AsyncIterable[int]):
+class AsyncSonyFlake:
     """Async wrapper for :class:`SonyFlake`.
 
     Usage:
@@ -57,10 +57,9 @@ class AsyncSonyFlake(Awaitable[int], AsyncIterable[int]):
                 break
     """
 
-    __slots__: ClassVar[tuple[str, ...]] = ("sf", "sleep")
-    sleep: AsyncSleep
+    __slots__ = ("sf", "sleep")
 
-    def __init__(self, sf: SonyFlake, sleep: AsyncSleep | None = None) -> None:
+    def __init__(self, sf: SonyFlake, sleep=None):
         """Initialize AsyncSonyFlake ID generator.
 
         Args:
@@ -76,7 +75,7 @@ def __init__(self, sf: SonyFlake, sleep: AsyncSleep | None = None) -> None:
         self.sf = sf
         self.sleep = sleep
 
-    async def __call__(self, n: int) -> list[int]:
+    async def __call__(self, n):
         """Generate multiple SonyFlake IDs at once.
 
         Roughly equivalent to ``[await asf for _ in range(n)]``, but more
@@ -96,7 +95,7 @@ async def __call__(self, n: int) -> list[int]:
 
         return ids
 
-    def __await__(self) -> Generator[Any, Any, int]:
+    def __await__(self):
         """Produce a SonyFlake ID."""
 
         id_, to_sleep = self.sf._raw(None)
@@ -105,12 +104,12 @@ def __await__(self) -> Generator[Any, Any, int]:
 
         return id_
 
-    def __aiter__(self) -> AsyncIterator[int]:
+    def __aiter__(self):
         """Return an infinite SonyFlake ID async iterator."""
 
         return self._gen().__aiter__()
 
-    async def _gen(self) -> AsyncIterator[int]:
+    async def _gen(self):
         """Infinite SonyFlake ID async generator."""
 
         while True:
diff --git a/src/sonyflake_turbo/__init__.pyi b/src/sonyflake_turbo/__init__.pyi
@@ -0,0 +1,64 @@
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable
+from typing import (
+    Any,
+    Callable,
+    ClassVar,
+    Generator,
+    TypeAlias,
+)
+
+from ._sonyflake import (
+    SONYFLAKE_EPOCH,
+    SONYFLAKE_MACHINE_ID_BITS,
+    SONYFLAKE_MACHINE_ID_MAX,
+    SONYFLAKE_MACHINE_ID_OFFSET,
+    SONYFLAKE_SEQUENCE_BITS,
+    SONYFLAKE_SEQUENCE_MAX,
+    SONYFLAKE_TIME_OFFSET,
+    MachineIDLCG,
+    SonyFlake,
+)
+
+__all__ = [
+    "SONYFLAKE_EPOCH",
+    "SONYFLAKE_MACHINE_ID_BITS",
+    "SONYFLAKE_MACHINE_ID_MAX",
+    "SONYFLAKE_MACHINE_ID_OFFSET",
+    "SONYFLAKE_SEQUENCE_BITS",
+    "SONYFLAKE_SEQUENCE_MAX",
+    "SONYFLAKE_TIME_OFFSET",
+    "AsyncSonyFlake",
+    "MachineIDLCG",
+    "SonyFlake",
+]
+
+AsyncSleep: TypeAlias = Callable[[float], Awaitable[None]]
+
+class AsyncSonyFlake(Awaitable[int], AsyncIterable[int]):
+    """Async wrapper for :class:`SonyFlake`.
+
+    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
+    """
+
+    __slots__: ClassVar[tuple[str, ...]]
+    sleep: AsyncSleep
+
+    def __init__(self, sf: SonyFlake, sleep: AsyncSleep | None = None) -> None: ...
+    async def __call__(self, n: int) -> list[int]: ...
+    def __await__(self) -> Generator[Any, Any, int]: ...
+    def __aiter__(self) -> AsyncIterator[int]: ...
+    async def _gen(self) -> AsyncIterator[int]: ...
diff --git a/src/sonyflake_turbo/pure.py b/src/sonyflake_turbo/pure.py
@@ -0,0 +1,176 @@
+from threading import Lock
+from time import sleep, time_ns
+from typing import Self, overload
+
+SONYFLAKE_EPOCH = 1409529600  # 2014-09-01 00:00:00 UTC
+SONYFLAKE_SEQUENCE_BITS = 8
+SONYFLAKE_SEQUENCE_COUNT = 1 << SONYFLAKE_SEQUENCE_BITS
+SONYFLAKE_SEQUENCE_MAX = SONYFLAKE_SEQUENCE_COUNT - 1
+SONYFLAKE_MACHINE_ID_BITS = 16
+SONYFLAKE_MACHINE_ID_MAX = (1 << SONYFLAKE_MACHINE_ID_BITS) - 1
+SONYFLAKE_MACHINE_ID_OFFSET = SONYFLAKE_SEQUENCE_BITS
+SONYFLAKE_TIME_OFFSET = SONYFLAKE_MACHINE_ID_BITS + SONYFLAKE_SEQUENCE_BITS
+
+
+def sf_to_ns(start_time: int, sf: int) -> int:
+    return (start_time + sf) * 10_000_000
+
+
+def ns_to_sf(start_time: int, ns: int) -> int:
+    return ns // 10_000_000 - start_time
+
+
+def compose(i: int, elapsed: int, machine_ids: list[int]) -> int:
+    t = elapsed << SONYFLAKE_TIME_OFFSET
+    m = machine_ids[i >> SONYFLAKE_SEQUENCE_BITS] << SONYFLAKE_MACHINE_ID_OFFSET
+    c = i & SONYFLAKE_SEQUENCE_MAX
+    return t | m | c
+
+
+def diff(start_time: int, elapsed: int, current_ns: int) -> float:
+    d = sf_to_ns(start_time, elapsed) - current_ns
+    if d > 0:
+        return d / 1_000_000_000
+    return 0
+
+
+def machine_id_lcg(x: int) -> int:
+    return (32309 * x + 13799) % 65536
+
+
+class SonyFlake:
+    __slots__ = ("_machine_ids", "_start_time", "_elapsed", "_max_i", "_i", "_lock")
+
+    _start_time: int
+    _elapsed: int
+    _machine_ids: list[int]
+    _max_i: int
+    _i: int
+    _lock: Lock
+
+    def __init__(self, *machine_ids: int, start_time: int = SONYFLAKE_EPOCH):
+        machine_ids_set = set(machine_ids)
+
+        if len(machine_ids_set) == 0:
+            raise ValueError("At least one machine ID must be provided")
+
+        if len(machine_ids_set) > 65536:
+            raise ValueError("Too many machine IDs, maximum is 65536")
+
+        if len(machine_ids_set) != len(machine_ids):
+            raise ValueError("Duplicate machine IDs are not allowed")
+
+        if not all(map(lambda x: isinstance(x, int), machine_ids_set)):
+            raise TypeError("Machine IDs must be integers")
+
+        if min(machine_ids_set) < 0 or max(machine_ids_set) > 0xFFFF:
+            raise ValueError("Machine IDs must be in range [0, 65535]")
+
+        if not isinstance(start_time, int):
+            raise TypeError("start_time must be an integer")
+
+        self._machine_ids = sorted(machine_ids_set)
+        self._max_i = len(machine_ids_set) * SONYFLAKE_SEQUENCE_COUNT
+        self._start_time = start_time * 100
+        self._elapsed = 0
+        self._i = 0
+        self._lock = Lock()
+
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> int:
+        id_, to_sleep = self._next()
+        if to_sleep > 0:
+            sleep(to_sleep)
+        return id_
+
+    def __call__(self, n: int) -> list[int]:
+        ids, to_sleep = self._next_n(n)
+        sleep(to_sleep)
+        return ids
+
+    def __repr__(self) -> str:
+        cls = self.__class__.__name__
+        machine_ids = ", ".join(map(str, sorted(self._machine_ids)))
+        return f"{cls}({machine_ids}, start_time={self._start_time // 100})"
+
+    @overload
+    def _raw(self, n: int) -> tuple[list[int], float]: ...
+    @overload
+    def _raw(self, n: None) -> tuple[int, float]: ...
+
+    def _raw(self, n: int | None) -> tuple[int | list[int], float]:
+        if n is None:
+            return self._next()
+        return self._next_n(n)
+
+    def _next(self) -> tuple[int, float]:
+        start_time = self._start_time
+        machine_ids = self._machine_ids
+
+        with self._lock:
+            current_ns = time_ns()
+            current = ns_to_sf(start_time, current_ns)
+            elapsed = self._elapsed
+
+            if elapsed < current:
+                self._elapsed = elapsed = current
+                self._i = i = 0
+            else:
+                self._i = i = (self._i + 1) % self._max_i
+
+                if i == 0:
+                    self._elapsed = elapsed = elapsed + 1
+
+        return (
+            compose(i, elapsed, machine_ids),
+            diff(start_time, elapsed, current_ns),
+        )
+
+    def _next_n(self, n: int) -> tuple[list[int], float]:
+        if n < 1:
+            return [], 0
+
+        start_time = self._start_time
+        machine_ids = self._machine_ids
+        max_i = self._max_i
+        result: list[int] = []
+
+        with self._lock:
+            current_ns = time_ns()
+            current = ns_to_sf(start_time, current_ns)
+            elapsed = self._elapsed
+
+            if elapsed < current:
+                elapsed = current
+                i = 0
+            else:
+                i = self._i
+
+            while len(result) < n:
+                result.append(compose(i, elapsed, machine_ids))
+
+                i = (i + 1) % max_i
+
+                if i == 0:
+                    elapsed += 1
+
+            self._elapsed = elapsed
+            self._i = i
+
+        return result, diff(start_time, elapsed, current_ns)
+
+
+class MachineIDLCG:
+    def __init__(self, seed: int = 0):
+        self._seed = seed
+        self._lock = Lock()
+
+    def __next__(self) -> int:
+        with self._lock:
+            self._seed = x = machine_id_lcg(self._seed)
+            return x
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self._seed})"
diff --git a/tests/test_async.py b/tests/test_async.py
diff --git a/tests/test_sonyflake_turbo.py b/tests/test_sonyflake_turbo.py
