From c9a029efb94099b7cc8387093694cf661f4be7ba Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Thu, 9 Jan 2025 14:58:51 +0100
Subject: [PATCH] Allow passing Python timestamp function to undo manager
 (#214)

---
 python/pycrdt/_pycrdt.pyi |  2 +-
 python/pycrdt/_undo.py    | 12 +++++++++---
 src/undo.rs               | 23 +++++++++++++++++++++--
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/python/pycrdt/_pycrdt.pyi b/python/pycrdt/_pycrdt.pyi
index 105618a..9d8fd36 100644
--- a/python/pycrdt/_pycrdt.pyi
+++ b/python/pycrdt/_pycrdt.pyi
@@ -342,7 +342,7 @@ class XmlText:
 class UndoManager:
     """Undo manager."""
 
-    def __init__(self, doc: Doc, capture_timeout_millis: int = 500) -> None:
+    def __init__(self, doc: Doc, capture_timeout_millis, timestamp: Callable[[], int]) -> None:
         """Creates an undo manager."""
 
     def expand_scope(self, scope: Text | Array | Map) -> None:
diff --git a/python/pycrdt/_undo.py b/python/pycrdt/_undo.py
index b50a877..99a1a2d 100644
--- a/python/pycrdt/_undo.py
+++ b/python/pycrdt/_undo.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from time import time_ns
+from typing import TYPE_CHECKING, Any, Callable
 
 from ._base import BaseType
 from ._pycrdt import (
@@ -15,6 +16,10 @@
     from ._doc import Doc
 
 
+def timestamp() -> int:
+    return time_ns() // 1_000_000
+
+
 class UndoManager:
     """
     The undo manager allows to perform undo/redo operations on shared types.
@@ -31,13 +36,14 @@ def __init__(
         doc: Doc | None = None,
         scopes: list[BaseType] = [],
         capture_timeout_millis: int = 500,
+        timestamp: Callable[[], int] = timestamp,
     ) -> None:
         """
         Args:
             doc: The document the undo manager will work with.
             scopes: A list of shared types the undo manager will work with.
             capture_timeout_millis: A time interval for grouping changes that will be undone/redone.
-
+            timestamp: A function that returns a timestamp as an integer number of milli-seconds.
 
         Raises:
             RuntimeError: UndoManager must be created with doc or scopes.
@@ -48,7 +54,7 @@ def __init__(
             doc = scopes[0].doc
         elif scopes:
             raise RuntimeError("UndoManager must be created with doc or scopes")
-        self._undo_manager = _UndoManager(doc._doc, capture_timeout_millis)
+        self._undo_manager = _UndoManager(doc._doc, capture_timeout_millis, timestamp)
         for scope in scopes:
             self.expand_scope(scope)
 
diff --git a/src/undo.rs b/src/undo.rs
index 694090d..a4921f7 100644
--- a/src/undo.rs
+++ b/src/undo.rs
@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+use std::sync::Arc;
 use pyo3::prelude::*;
 use pyo3::types::PyList;
 use pyo3::exceptions::PyRuntimeError;
@@ -8,11 +10,23 @@ use yrs::undo::{
     Options,
     StackItem as _StackItem,
 };
+use yrs::sync::{Clock, Timestamp};
 use crate::doc::Doc;
 use crate::text::Text;
 use crate::array::Array;
 use crate::map::Map;
 
+struct PythonClock {
+    timestamp: PyObject,
+}
+
+impl Clock for PythonClock {
+    fn now(&self) -> Timestamp {
+        Python::with_gil(|py| {
+            self.timestamp.call0(py).expect("Error getting timestamp").extract(py).expect("Could not convert timestamp to int")
+        })
+    }
+}
 
 #[pyclass(unsendable)]
 pub struct UndoManager {
@@ -22,8 +36,13 @@ pub struct UndoManager {
 #[pymethods]
 impl UndoManager {
     #[new]
-    fn new(doc: &Doc, capture_timeout_millis: u64) -> Self {
-        let mut options = Options::default();
+    fn new(doc: &Doc, capture_timeout_millis: u64, timestamp: PyObject) -> Self {
+        let mut options = Options {
+            capture_timeout_millis: 500,
+            tracked_origins: HashSet::new(),
+            capture_transaction: None,
+            timestamp: Arc::new(PythonClock {timestamp}),
+        };
         options.capture_timeout_millis = capture_timeout_millis;
         let undo_manager = _UndoManager::with_options(&doc.doc, options);
         UndoManager { undo_manager }