support mock.patch.object on HasTraits

tests/test_traitlets.py

@@ -12,7 +12,7 @@
 import pickle
 import re
 import typing as t
-from unittest import TestCase
+from unittest import TestCase, mock
 
 import pytest
 
@@ -3199,3 +3199,19 @@ def test_all_attribute():
     for name in traitlets.__all__:
         if name not in names:
             raise ValueError(f"{name} should be removed from __all__")
+
+
+def test_mock_patch():
+    class A(HasTraits):
+        trait = Unicode(default_value="default")
+
+    a = A()
+    # not set, restores default
+    with mock.patch.object(a, "trait", "patch"):
+        assert a.trait == "patch"
+    assert a.trait == "default"
+    # set, restores before state
+    a.trait = "set"
+    with mock.patch.object(a, "trait", "patch"):
+        assert a.trait == "patch"
+    assert a.trait == "set"

traitlets/traitlets.py

@@ -537,6 +537,9 @@ def instance_init(self, obj: t.Any) -> None:
     V = TypeVar("V")
 
 
+_DELETED = object()
+
+
 # We use a type for the getter (G) and setter (G) because we allow
 # for traits to cast (for instance CInt will use G=int, S=t.Any)
 class TraitType(BaseDescriptor, t.Generic[G, S]):
@@ -726,6 +729,10 @@ def __get__(self, obj: HasTraits | None, cls: type[t.Any]) -> Self | G:
         if obj is None:
             return self
         else:
+            if obj._trait_values.get(self.name, None) is _DELETED:
+                # if _DELETED sentinel is set, behave as if attribute is not set
+                # otherwise delattr does weird things
+                raise AttributeError(self.name)
             return self.get(obj, cls)  # type:ignore[return-value]
 
     def set(self, obj: HasTraits, value: S) -> None:
@@ -757,6 +764,12 @@ def __set__(self, obj: HasTraits, value: S) -> None:
             raise TraitError(f'The "{self.name}" trait is read-only.')
         self.set(obj, value)
 
+    def __delete__(self, obj: HasTraits) -> None:
+        """
+        delattr stores a sentinel so `hasattr` returns False
+        """
+        obj._trait_values[self.name] = _DELETED
+
     def _validate(self, obj: t.Any, value: t.Any) -> G | None:
         if value is None and self.allow_none:
             return value