model_copy doesn't reset ReferenceMixin._ref weakrefs

[jmuhlich]

I noticed some strange behavior with someone else's code that uses copy.deepcopy to clone an OME object. I tried the pydantic-provided model_copy(deep=True) method and had the same bad results. The problem seems to be that the ReferenceMixin._ref weakrefs in the copy still reference the original objects rather than the deep-copied instances. This leads to false sharing of those objects as well as refs that unexpectedly become "dead" depending on the lifetime of the original objects. I was able to make it work with OME(**original.dict()) (ugly) or calling the private OME._link_refs (poor form). I thought our OME.__setstate__ would address this but apparently not. There's probably a relevant pydantic hook where we need to call _link_refs, but I'm lost in the codebase after the refactoring and move to pydantic 2. 😅

Here's a code snippet that demonstrates the problems. In the real-world case I observed there was a function call boundary involved but I've used del to simulate the object lifecycle issues instead.

import copy
from ome_types.model import OME, Instrument, AnnotationRef, CommentAnnotation

aref = AnnotationRef(id=1)
ome = OME(
    instruments=[Instrument(annotation_refs=[aref])],
    structured_annotations=[CommentAnnotation(id=1, value="test")],
)

ome2 = ome.model_copy(deep=True)
print(
    "ref alias error - pydantic model_copy: ",
    ome2.instruments[0].annotation_refs[0].ref is aref.ref,
)
ome3 = copy.deepcopy(ome)
print(
    "ref alias error - copy.deepcopy:       ",
    ome3.instruments[0].annotation_refs[0].ref is aref.ref,
)
ome4 = OME(**ome.dict())
print(
    "ref alias error - copy via constructor:",
    ome4.instruments[0].annotation_refs[0].ref is aref.ref,
)

del ome, aref
print()
print(
    "weakref dead - pydantic model_copy:    ",
    ome2.instruments[0].annotation_refs[0].ref is None
)
print(
    "weakref dead - copy.deepcopy:          ",
    ome3.instruments[0].annotation_refs[0].ref is None
)
print(
    "weakref dead - copy via constructor:   ",
    ome4.instruments[0].annotation_refs[0].ref is None
)

[jmuhlich]

Note that you can also observe the issue by verifying whether omeX.instruments[0].annotation_refs[0].ref is omeX.structured_annotations[0] (should be True, but is False for the model_copy and deepcopy cases).

[tlambert03]

can you check if just updating the OMEMixin as follows does what you expect?

# src/ome_types/_mixins/_ome.py
class OMEMixin:
    ...
    def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self:
        copy = super().__deepcopy__(memo)
        copy._link_refs()
        return copy

with that change I get:

ref alias error - pydantic model_copy:  False
ref alias error - copy.deepcopy:        False
ref alias error - copy via constructor: False

weakref dead - pydantic model_copy:     False
weakref dead - copy.deepcopy:           False
weakref dead - copy via constructor:    False