Context change hook

tests/test_hybrid_class.py

@@ -405,3 +405,96 @@ def c(self):
     assert all(td3.a == td2.a)
     assert td3.b == td2.b
     assert td3.c == td2.c
+
+
+def test_on_context_change():
+    class TestClass(xo.HybridClass):
+        _xofields = {
+            "value": xo.Float64,
+        }
+
+        def __init__(self, **kwargs):
+            self.context_changes = []
+            super().__init__(**kwargs)
+
+        def _on_context_change(self, old_context, new_context):
+            self.context_changes.append((old_context, new_context))
+
+    context0 = xo.ContextCpu()
+    context1 = xo.ContextCpu(omp_num_threads="auto")
+
+    # Initial construction
+    obj = TestClass(value=3.0, _context=context0)
+
+    assert obj._context is context0
+    assert obj.context_changes == [
+        (None, context0),
+    ]
+
+    # Move to a different context
+    obj.move(_context=context1)
+
+    assert obj._context is context1
+    assert obj.context_changes == [
+        (None, context0),
+        (context0, context1),
+    ]
+
+    # Moving within the same context should not trigger the hook
+    buffer1 = context1.new_buffer()
+    obj.move(_buffer=buffer1)
+
+    assert obj._context is context1
+    assert obj.context_changes == [
+        (None, context0),
+        (context0, context1),
+    ]
+
+    # Copy to another context. The hook must run on the copy,
+    # without modifying the source object's history.
+    copied = obj.copy(_context=context0)
+
+    assert copied is not obj
+    assert copied._context is context0
+
+    assert obj.context_changes == [
+        (None, context0),
+        (context0, context1),
+    ]
+
+    assert copied.context_changes == [
+        (None, context0),
+        (context0, context1),
+        (context1, context0),
+    ]
+
+    # A copy in the same context should not add another event.
+    copied_same_context = obj.copy()
+
+    assert copied_same_context._context is context1
+    assert copied_same_context.context_changes == [
+        (None, context0),
+        (context0, context1),
+    ]
+
+
+def test_on_context_change_with_direct_xoinitialize():
+    class TestClass(xo.HybridClass):
+        _xofields = {
+            "value": xo.Float64,
+        }
+
+        def __init__(self, **kwargs):
+            self.context_changes = []
+            self.xoinitialize(**kwargs)
+
+        def _on_context_change(self, old_context, new_context):
+            self.context_changes.append((old_context, new_context))
+
+    context = xo.ContextCpu()
+
+    obj = TestClass(value=3.0, _context=context)
+
+    assert obj.context_changes == [
+        (None, context),
+    ]

xobjects/hybrid_class.py

@@ -3,6 +3,7 @@
 # Copyright (c) CERN, 2021.                   #
 # ########################################### #
 
+from hashlib import new
 import json
 from inspect import isclass
 
@@ -214,11 +215,17 @@ def move(self, _context=None, _buffer=None, _offset=None):
                 "references to other objects."
             )
 
+        old_context = self._context
+
         self._xobject = self._xobject.__class__(
             self._xobject, _context=_context, _buffer=_buffer, _offset=_offset
         )
         self._reinit_from_xobject(_xobject=self._xobject)
 
+        new_context = self._context
+        if new_context is not old_context:
+            self._on_context_change(old_context, new_context)
+
     @property
     def _move_to(self):
         raise NameError("`_move_to` has been removed. Use `move` instead.")
@@ -255,6 +262,7 @@ def xoinitialize(self, _xobject=None, _kwargs_name_check=True, **kwargs):
 
         if _xobject is not None:
             self._reinit_from_xobject(_xobject=_xobject)
+            self._on_context_change(None, self._context)
             return
 
         # Handle dressed inputs
@@ -276,6 +284,8 @@ def xoinitialize(self, _xobject=None, _kwargs_name_check=True, **kwargs):
         # (for example in case object is initialized from dict)
         self._reinit_from_xobject(_xobject=self._xobject)
 
+        self._on_context_change(None, self._context)
+
     def __init__(self, _xobject=None, **kwargs):
         self.xoinitialize(_xobject=_xobject, **kwargs)
 
@@ -336,8 +346,9 @@ def from_dict(cls, dct, _context=None, _buffer=None, _offset=None):
         )
 
     def copy(self, _context=None, _buffer=None, _offset=None):
+        old_context = self._xobject._buffer.context
         if _context is None and _buffer is None:
-            _context = self._xobject._buffer.context
+            _context = old_context
         # This makes a copy of the xobject
         new_xobject = self._XoStruct(
             self._xobject, _context=_context, _buffer=_buffer, _offset=_offset
@@ -350,6 +361,10 @@ def copy(self, _context=None, _buffer=None, _offset=None):
             if hasattr(vv, "copy"):
                 new.__dict__[kk] = vv.copy()
         new._xobject = new_xobject
+
+        new_context = new._context
+        if new_context is not old_context:
+            new._on_context_change(old_context, new_context)
         return new
 
     @property
@@ -372,6 +387,7 @@ def __setstate__(self, state):
             buffer=state[0], offset=state[1]
         )
         self._reinit_from_xobject(_xobject=self._xobject)
+        self._on_context_change(None, self._context)
 
     @property
     def XoStruct(self):
@@ -418,3 +434,18 @@ def __repr__(self):
                 vvrepr = repr(vv)
             args.append(f"{fname}={vvrepr}")
         return f'{type(self).__name__}({", ".join(args)})'
+
+    def _on_context_change(self, old_context, new_context):
+        """Hook called after the backing xobject changes context.
+        It may inspect the backing Xobject/context, but should not assume that
+        arbitrary pure-Python subclass state has already been initialised.
+
+        Parameters
+        ----------
+        old_context : XContext or None
+            Previous context. None when the object is first initialised
+            or reconstructed.
+        new_context : XContext
+            Context containing the object's current backing xobject.
+        """
+        pass