diff --git a/tests/test_hybrid_class.py b/tests/test_hybrid_class.py index f4b0676..b388c10 100644 --- a/tests/test_hybrid_class.py +++ b/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), + ] diff --git a/xobjects/hybrid_class.py b/xobjects/hybrid_class.py index f83086c..a345e89 100644 --- a/xobjects/hybrid_class.py +++ b/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