Add optional nanos argument to uuid6 and uuid7 (#51)

Co-authored-by: Amin Alaee <[email protected]>
aminalaee · Jun 14, 2024 · 8c9391b · 8c9391b
1 parent 4838ddd
commit 8c9391b
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 6 deletions.
diff --git a/python/uuid_utils/__init__.pyi b/python/uuid_utils/__init__.pyi
@@ -151,14 +151,16 @@ def uuid5(namespace: UUID, name: str) -> UUID:
     """Generate a UUID from the SHA-1 hash of a namespace UUID and a name."""
     ...
 
-def uuid6(node: _Int | None = None, timestamp: _Int | None = None) -> UUID:
+def uuid6(
+    node: _Int | None = None, timestamp: _Int | None = None, nanos: _Int | None = None
+) -> UUID:
     """Generate a version 6 UUID using the given timestamp and a host ID.
     This is similar to version 1 UUIDs,
     except that it is lexicographically sortable by timestamp.
     """
     ...
 
-def uuid7(timestamp: _Int | None = None) -> UUID:
+def uuid7(timestamp: _Int | None = None, nanos: _Int | None = None) -> UUID:
     """Generate a version 7 UUID using a time value and random bytes."""
     ...
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -338,7 +338,7 @@ fn uuid5(namespace: &UUID, name: &str) -> PyResult<UUID> {
 }
 
 #[pyfunction]
-fn uuid6(node: Option<u64>, timestamp: Option<u64>) -> PyResult<UUID> {
+fn uuid6(node: Option<u64>, timestamp: Option<u64>, nanos: Option<u32>) -> PyResult<UUID> {
     let node = match node {
         Some(node) => node.to_ne_bytes(),
         None => _getnode().to_ne_bytes(),
@@ -347,7 +347,7 @@ fn uuid6(node: Option<u64>, timestamp: Option<u64>) -> PyResult<UUID> {
 
     let uuid = match timestamp {
         Some(timestamp) => {
-            let timestamp = Timestamp::from_unix(&Context::new_random(), timestamp, 0);
+            let timestamp = Timestamp::from_unix(&Context::new_random(), timestamp, nanos.unwrap_or(0));
             return Ok(UUID {
                 uuid: Uuid::new_v6(timestamp, node),
             });
@@ -358,10 +358,10 @@ fn uuid6(node: Option<u64>, timestamp: Option<u64>) -> PyResult<UUID> {
 }
 
 #[pyfunction]
-fn uuid7(timestamp: Option<u64>) -> PyResult<UUID> {
+fn uuid7(timestamp: Option<u64>, nanos: Option<u32>) -> PyResult<UUID> {
     let uuid = match timestamp {
         Some(timestamp) => {
-            let timestamp = Timestamp::from_unix(&Context::new_random(), timestamp, 0);
+            let timestamp = Timestamp::from_unix(&Context::new_random(), timestamp, nanos.unwrap_or(0));
             return Ok(UUID {
                 uuid: Uuid::new_v7(timestamp),
             });

diff --git a/tests/test_uuid.py b/tests/test_uuid.py
@@ -1,6 +1,7 @@
 import copy
 import pickle
 import sys
+from datetime import datetime
 from uuid import UUID, getnode
 
 import pytest
@@ -92,6 +93,9 @@ def test_uuid6() -> None:
     uuid = uuid_utils.uuid6(getnode(), 1679665408)
     assert isinstance(uuid, uuid_utils.UUID)
 
+    uuid = uuid_utils.uuid6(getnode(), 1679665408, 123)
+    assert isinstance(uuid, uuid_utils.UUID)
+
     uuid = uuid_utils.uuid6()
     assert isinstance(uuid, uuid_utils.UUID)
 
@@ -100,9 +104,18 @@ def test_uuid7() -> None:
     uuid = uuid_utils.uuid7(1679665408)
     assert isinstance(uuid, uuid_utils.UUID)
 
+    uuid = uuid_utils.uuid7(1679665408, 999)
+    assert isinstance(uuid, uuid_utils.UUID)
+
     uuid = uuid_utils.uuid7()
     assert isinstance(uuid, uuid_utils.UUID)
 
+    ts = datetime(
+        year=2024, month=1, day=2, hour=3, minute=4, second=5, microsecond=123000
+    )
+    uuid = uuid_utils.uuid7(int(ts.timestamp()), ts.microsecond * 1_000)
+    assert uuid.timestamp == int(ts.timestamp() * 1000)
+
 
 def test_uuid8() -> None:
     uuid = uuid_utils.uuid8(b"1234567812345678")