[freddieknets]

Description

Add context-change lifecycle hook to HybridClass

This PR introduces an _on_context_change(old_context, new_context) lifecycle hook for HybridClass.

The hook allows classes backed by Xobjects to react when their backing Xobject is created in, moved to, or copied to a different execution context.

Motivation

Until now, subclasses that need context-specific setup have had to implement this separately in their constructor, move(), and copy() methods.

This is both repetitive and easy to get wrong, since moving or copying an Xobject does not rerun the subclass constructor.

A concrete example is Xcoll's EverestCrystal. CUDA tracking of crystals requires a larger per-thread CUDA stack. Initially this setup was performed only in EverestCrystal.__init__():

set_crystal_stack_limit(self._context)

This worked when a crystal was constructed directly on ContextCupy, but not when it was first constructed on CPU and subsequently moved to CUDA as part of Line.build_tracker(): HybridClass.move() replaces the backing Xobject without rerunning __init__().

The same issue applies to copy(_context=...).

More generally, subclasses may need to initialise or validate resources whenever they enter a particular context. A common lifecycle hook provides one place to implement such behaviour.

API

Subclasses can implement:

def _on_context_change(self, old_context, new_context):
    ...

where:

• old_context is None during initial construction;
• new_context is the context of the current backing Xobject;
• on a genuine context transition, both the old and new contexts are provided.

The base implementation is a no-op.

For example:

class MyClass(xo.HybridClass):

    def _on_context_change(self, old_context, new_context):
        if isinstance(new_context, xo.ContextCupy):
            initialise_cuda_resources(new_context)

Behaviour

The hook is invoked:

• after initial Xobject construction, as

  _on_context_change(None, new_context)

• after move() when the context changes;

• on the newly created object after copy() when the destination context differs from the source context.

It is not invoked when only the backing buffer changes while remaining in the same context.

Initialisation is hooked through xoinitialize() rather than only through HybridClass.__init__(), since subclasses are allowed to call xoinitialize() directly.

Example

context_cpu = xo.ContextCpu()
context_gpu = xo.ContextCupy()

obj = MyClass(_context=context_cpu)
# _on_context_change(None, context_cpu)

obj.move(_context=context_gpu)
# _on_context_change(context_cpu, context_gpu)

obj_gpu2 = obj.copy(_context=context_cpu)
# obj_gpu2._on_context_change(context_gpu, context_cpu)

Tests

The added tests verify that:

• the hook is called during initial construction;
• move() reports the correct source and destination contexts;
• moving to a different buffer in the same context does not trigger the hook;
• copy() invokes the hook on the newly created object without modifying the source object;
• __setstate__() invokes the hook
• copying within the same context does not trigger an additional context-change event;
• subclasses calling xoinitialize() directly receive the initial lifecycle callback.

This keeps context-dependent setup in one well-defined place and avoids subclasses having to duplicate lifecycle handling across construction, move, and copy operations.

Checklist

Mandatory:

• I have added tests to cover my changes
• All the tests are passing, including my new ones
• I described my changes in this PR description

Optional:

• The code I wrote follows good style practices (see PEP 8 and PEP 20).
• I have updated the docs in relation to my changes, if applicable
• I have tested also GPU contexts