uid test

src/jobflow/utils/uid.py

@@ -0,0 +1,98 @@
+"""Tools for generating UUIDs."""
+from __future__ import annotations
+
+from uuid import UUID
+
+import ulid
+
+
+def suid(id_type: str | None = None) -> str:
+    """Generate a string UUID (universally unique identifier).
+
+    Since the timestamp of the IDs are important for sorting,
+    only id types that include a timestamp are supported.
+
+    Parameters
+    ----------
+    uuid_type:
+        The type of UUID to generate.
+        In the future, ``uuid7`` and ``ulid`` may be supported.
+
+    Returns
+    -------
+    str
+        A UUID.
+    """
+    import uuid
+
+    from jobflow import SETTINGS
+
+    if id_type is None:
+        id_type = SETTINGS.UID_TYPE
+
+    funcs = {
+        "uuid1": uuid.uuid1,
+        "uuid4": uuid.uuid4,
+        "ulid": ulid.ULID,
+    }
+    if id_type not in funcs:
+        raise ValueError(f"UUID type {id_type} not supported.")
+    return str(funcs[id_type]())
+
+
+def get_timestamp_from_uid(uid: str) -> float:
+    """
+    Get the time that a UID was generated.
+
+    Parameters
+    ----------
+    uuid
+        A UUID.
+
+    Returns
+    -------
+    float
+        The time stamp from the UUID.
+    """
+    id_type = _get_id_type(uid)
+    if id_type == "uuid4":
+        raise ValueError(
+            "UUID4 is randomly generated and not associated with a time stamp."
+        )
+    funcs = {
+        "uuid1": lambda uuid: (UUID(uuid).time - 0x01B21DD213814000) / 1e7,
+        "ulid": lambda uuid: ulid.ULID.from_str(uuid).timestamp,
+    }
+    return funcs[id_type](uid)
+
+
+def _get_id_type(uid: str) -> str:
+    """
+    Get the type of a UUID.
+
+    Parameters
+    ----------
+    uuid
+        A UUID.
+
+    Returns
+    -------
+    str
+        The type of the UUID.
+    """
+    try:
+        version = UUID(uid).version
+        return {
+            1: "uuid1",
+            4: "uuid4",
+        }[version]
+    except ValueError:
+        pass
+
+    try:
+        ulid.ULID.from_str(uid)
+        return "ulid"
+    except ValueError:
+        pass
+
+    raise ValueError(f"ID type for {uid} not recognized.")

tests/utils/test_uid.py

@@ -0,0 +1,38 @@
+import pytest
+
+
+def test_uid():
+    from uuid import UUID
+
+    from ulid import ULID
+
+    from jobflow.utils.uid import get_timestamp_from_uid, suid
+
+    uuid = suid("uuid1")
+    assert UUID(uuid).version == 1
+    t1 = get_timestamp_from_uid(uuid)
+    assert isinstance(t1, float)
+
+    uuid = suid("uuid4")
+    assert UUID(uuid).version == 4
+
+    with pytest.raises(
+        ValueError,
+        match="UUID4 is randomly generated and not associated with a time stamp.",
+    ):
+        get_timestamp_from_uid(uuid)
+
+    ulid = suid("ulid")
+    assert ULID.from_str(ulid)
+    t2 = get_timestamp_from_uid(ulid)
+    assert isinstance(t2, float)
+
+    with pytest.raises(ValueError, match="UUID type uuid2 not supported."):
+        suid("uuid2")
+
+    with pytest.raises(ValueError, match="ID type for FAKEUUID not recognized."):
+        get_timestamp_from_uid("FAKEUUID")
+
+    default_uid = suid()
+    assert UUID(default_uid).version == 4
+    # assert len(ULID.from_str(default_uid).hex) == 32 # uncomment when ulid is default