From f9d1d4db075b7afcd95ea2abf91d7eb882fb3ed4 Mon Sep 17 00:00:00 2001
From: balex89 <balex89@gmail.com>
Date: Tue, 21 Jan 2025 21:53:36 +0300
Subject: [PATCH] Copyable classes with bindable properties

---
 nicegui/binding.py    | 48 ++++++++++++++++++++++++++++++++++++++++---
 tests/test_binding.py | 24 +++++++++++++++++++++-
 2 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/nicegui/binding.py b/nicegui/binding.py
index d047d958b..c8dc06e38 100644
--- a/nicegui/binding.py
+++ b/nicegui/binding.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import copyreg
 import dataclasses
 import time
 from collections import defaultdict
@@ -33,10 +34,11 @@
 
 bindings: DefaultDict[Tuple[int, str], List] = defaultdict(list)
 bindable_properties: Dict[Tuple[int, str], Any] = {}
+copyable_classes: Set[Type] = set()
 active_links: List[Tuple[Any, str, Any, str, Callable[[Any], Any]]] = []
 
-T = TypeVar('T', bound=type)
-
+TC = TypeVar('TC', bound=type)
+T = TypeVar('T')
 
 def _has_attribute(obj: Union[object, Mapping], name: str) -> Any:
     if isinstance(obj, Mapping):
@@ -169,6 +171,8 @@ def __get__(self, owner: Any, _=None) -> Any:
 
     def __set__(self, owner: Any, value: Any) -> None:
         has_attr = hasattr(owner, '___' + self.name)
+        if not has_attr and type(owner) not in copyable_classes:
+            _make_copyable(type(owner))
         value_changed = has_attr and getattr(owner, '___' + self.name) != value
         if has_attr and not value_changed:
             return
@@ -214,7 +218,7 @@ def reset() -> None:
 
 
 @dataclass_transform()
-def bindable_dataclass(cls: Optional[T] = None, /, *,
+def bindable_dataclass(cls: Optional[TC] = None, /, *,
                        bindable_fields: Optional[Iterable[str]] = None,
                        **kwargs: Any) -> Union[Type[DataclassInstance], IdentityFunction]:
     """A decorator that transforms a class into a dataclass with bindable fields.
@@ -252,3 +256,41 @@ def wrap(cls_):
         bindable_property.__set_name__(dataclass, field_name)
         setattr(dataclass, field_name, bindable_property)
     return dataclass
+
+
+def _register_bindables(original_obj: T, copy_obj: T) -> None:
+    """Ensure BindableProperties of an object copy are registered correctly.
+
+    :param original_obj: The object that was copied.
+    :param original_obj: The object copy.
+    """
+    for attr_name in dir(original_obj):
+        if (id(original_obj), attr_name) in bindable_properties:
+            bindable_properties[(id(copy_obj), attr_name)] = copy_obj
+
+
+def _register_bindables_pickle_function(obj: T) -> Tuple[Callable[..., T], Tuple[Any, ...]]:
+    """Construct the "reduce tuple" of an object with a wrapped pickle function, that registers bindable attributes"
+
+    :param obj: The object to be reduced.
+    """
+    reduced = obj.__reduce__()
+    if isinstance(reduced, str):
+        raise ValueError('Unexpected __reduce__() return type: str')
+    creator = reduced[0]
+
+    def creator_with_hook(*args, **kwargs) -> T:
+        obj_copy = creator(*args, **kwargs)
+        _register_bindables(obj, obj_copy)
+        return obj_copy
+
+    return (creator_with_hook,) + reduced[1:]
+
+
+def _make_copyable(cls: Type[T]) -> None:
+    """Modify the way `copy` module handles class instances so that `BindableProperty` attributes preserve bindability.
+
+    :param cls: The class to modify.
+    """
+    copyreg.pickle(cls, _register_bindables_pickle_function)
+    copyable_classes.add(cls)
diff --git a/tests/test_binding.py b/tests/test_binding.py
index 2c0ad2e36..d806131cf 100644
--- a/tests/test_binding.py
+++ b/tests/test_binding.py
@@ -1,9 +1,10 @@
+import copy
 from typing import Dict, Optional, Tuple
 
 from selenium.webdriver.common.keys import Keys
 
 from nicegui import binding, ui
-from nicegui.testing import Screen
+from nicegui.testing import Screen, User
 
 
 def test_ui_select_with_tuple_as_key(screen: Screen):
@@ -125,3 +126,24 @@ class TestClass:
     assert len(binding.bindings) == 2
     assert len(binding.active_links) == 1
     assert binding.active_links[0][1] == 'not_bindable'
+
+
+def test_copy_instance_with_bindable_property(user: User):
+    class TestClass:
+        x = binding.BindableProperty()
+        y = binding.BindableProperty()
+
+        def __init__(self):
+            self.x = 1
+            self.y = 2
+
+    original = TestClass()
+    duplicate = copy.copy(original)
+
+    ui.number().bind_value_from(original, 'x')
+    ui.number().bind_value_from(original, 'y')
+    ui.number().bind_value_from(duplicate, 'x')
+    ui.number().bind_value_from(duplicate, 'y')
+
+    assert len(binding.bindings) == 4
+    assert len(binding.active_links) == 0