Python UUIDv7

uuidv7.py

import datetime
import functools
import itertools
import operator
import os
import time
import uuid

import rich


class _GlobalCounter:
    def __init__(self):
        self._incs = itertools.count(1)
        self._last_value = 0

    def next(self):
        if self._last_value > 2**14 - 2:
            self._incs = itertools.count(1)
        value = next(self._incs)
        self._last_value = value
        return value


_global_counter = _GlobalCounter()


def _print_bytes(*numbers: int):
    console = rich.get_console()
    breaks = [(0, 48), (48, 52), (52, 64), (64, 66), (66, 114), (114, 128)]
    colors = ["magenta", "cyan", "green", "yellow", "blue", "red"]
    for n, (start, end), color in zip(numbers, breaks, colors, strict=False):
        # format the slice using color leave the rest as dim
        string = f"{n:0128b}"
        prefix = string[:start]
        colored = string[start:end]
        suffix = string[end:]
        final_string = f"[dim]{prefix}[/][{color}]{colored}[/][dim]{suffix}[/]"
        if prefix:
            assert prefix == "0" * len(prefix)
        if suffix:
            assert suffix == "0" * len(suffix)
        console.print(final_string, highlight=False)

    final_number = functools.reduce(operator.or_, numbers)
    final_string = f"{final_number:0128b}"
    for (start, end), color in zip(breaks, colors, strict=False):
        console.print(final_string[start:end], style=color, highlight=False, end="")
    console.print()


def uuid7(unix_ts_ms: int | datetime.datetime | None = None) -> uuid.UUID:
    """
     UUID v7, following the proposed extension to RFC4122 described in
     https://www.ietf.org/id/draft-peabody-dispatch-new-uuid-format-02.html.
     All representations (string, byte array, int) sort chronologically,
     with a potential time resolution of 50ns (if the system clock
     supports this).

     Parameters
     ----------
     unix_ts_ms
        Optional integer with the whole number of milliseconds since Unix epoch,
        or a datetime object. If not given, the current system time is used.

     Implementation notes
     --------------------
     RFC Draft: https://datatracker.ietf.org/doc/rfc9562/

     0                   1                   2                   3
     0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
    +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    |                           unix_ts_ms                          |
    +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    |          unix_ts_ms           |  ver  |     micro_seconds     |
    +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    |var|                        rand                               |
    +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    |                               |            counter            |
    +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

    unix_ts_ms:
       48-bit big-endian unsigned number of the Unix Epoch timestamp in
       milliseconds as per Section 6.1. Occupies bits 0 through 47

    ver:
       The 4-bit version field as defined by Section 4.2, set to 0b0111
       (7).  Occupies bits 48 through 51.

    sub_ms:
       12 bits big-endian unsigned number representing the sub-millisecond
       portion of the timestamp. Occupies bits 52 through 63.

    var:
       The 2-bit variant field as defined by Section 4.1, set to 0b10.
       Occupies bits 64 and 65.

    rand:
       48 bits of pseudorandom data to provide uniqueness as per Section 6.9.
       Occupies bits 66 through 113.

    counter:
       14 bits of a monotonically increasing counter to provide uniqueness
       Occupies bits 114 through 127.
    """
    if unix_ts_ms is None:
        unix_ts_ms, _rest = divmod(time.time_ns(), 1_000_000)
        micro_seconds = _rest // 1_000
    elif isinstance(unix_ts_ms, int):
        micro_seconds = 0
    elif isinstance(unix_ts_ms, datetime.datetime):
        ms, _rest = divmod(unix_ts_ms.timestamp() * 1_000, 1)
        unix_ts_ms, micro_seconds = int(ms), int(_rest * 1_000)
    else:
        msg = "Invalid type for unix_ts_ms"
        raise AssertionError(msg)

    uuid_int = (
        unix_ts_ms << 80  # 48 bits for unix_ts_ms
        | 0b0111 << 76  # 4 bits for version (7)
        | micro_seconds << 64  # 12 bits for sub_ms
        | 0b10 << 62  # 2 bits for variant
        | int.from_bytes(os.urandom(6), "big") << 14  # 48 random bits
        | _global_counter.next()  # 16 bits for monotonic counter
    )
    return uuid.UUID(int=uuid_int)