Clean up the data_item.display_data_channel associations to fix issue…

… where display item not added to document closes properly.

nion/swift/model/DataItem.py

@@ -250,7 +250,7 @@ def __init__(self, data: typing.Optional[_ImageDataType] = None, item_uuid: typi
         self.__pending_xdata: typing.Optional[DataAndMetadata.DataAndMetadata] = None
         self.__pending_queue: typing.List[typing.Tuple[DataAndMetadata.DataAndMetadata, typing.Sequence[slice], typing.Sequence[slice], DataAndMetadata.DataMetadata]] = list()
         self.__content_changed = False
-        self.__display_data_channel_refs = set()  # type: ignore  # display data channels referencing this data item
+        self.__display_data_channel_refs = set[weakref.ReferenceType["DisplayItem.DisplayDataChannel"]]()
         if data is not None:
             data_and_metadata = DataAndMetadata.DataAndMetadata.from_data(data, timezone=self.timezone, timezone_offset=self.timezone_offset)
             self.increment_data_ref_count()
@@ -463,7 +463,11 @@ def item_specifier(self) -> Persistence.PersistentObjectSpecifier:
         return Persistence.PersistentObjectSpecifier(self.uuid)
 
     def add_display_data_channel(self, display_data_channel: DisplayItem.DisplayDataChannel) -> None:
-        """Add a display data channel referencing this data item."""
+        """Add a display data channel referencing this data item.
+
+        This is used as an optimization to be able to easily access the display data channels for a data item.
+        Think of it as an "index" in the database sense.
+        """
         self.__display_data_channel_refs.add(weakref.ref(display_data_channel))
         self.notify_add_item("display_data_channels", display_data_channel)
 
@@ -475,7 +479,12 @@ def remove_display_data_channel(self, display_data_channel: DisplayItem.DisplayD
     @property
     def display_data_channels(self) -> typing.Set[DisplayItem.DisplayDataChannel]:
         """Return the list of display data channels referencing this data item."""
-        return {display_data_channel_ref() for display_data_channel_ref in self.__display_data_channel_refs}
+        display_data_channels = set["DisplayItem.DisplayDataChannel"]()
+        for display_data_channel_ref in self.__display_data_channel_refs:
+            display_data_channel = display_data_channel_ref()
+            assert display_data_channel  # should never be None
+            display_data_channels.add(display_data_channel)
+        return display_data_channels
 
     @property
     def in_transaction_state(self) -> bool:

nion/swift/model/DisplayItem.py

@@ -900,6 +900,8 @@ def __init__(self, data_item: typing.Optional[DataItem.DataItem] = None) -> None
 
         self.__last_data_item: typing.Optional[DataItem.DataItem] = None
 
+        self.__is_data_item_connected = False
+
         def connect_data_item(data_item_: typing.Optional[Persistence.PersistentObject]) -> None:
             data_item = typing.cast(DataItem.DataItem, data_item_)
             self.__disconnect_data_item_events()
@@ -911,6 +913,7 @@ def connect_data_item(data_item_: typing.Optional[Persistence.PersistentObject])
             # tell the data item that this display data channel is referencing it
             if data_item:
                 data_item.add_display_data_channel(self)
+                self.__is_data_item_connected = True
             if self.__data_item:
                 for _ in range(self.__display_ref_count):
                     self.__data_item.increment_data_ref_count()
@@ -926,6 +929,7 @@ def disconnect_data_item(data_item_: typing.Optional[Persistence.PersistentObjec
             # tell the data item that this display data channel is no longer referencing it
             if data_item:
                 data_item.remove_display_data_channel(self)
+                self.__is_data_item_connected = False
             self.notify_property_changed("data_item")
 
         self.__data_item_reference.on_item_registered = connect_data_item
@@ -935,12 +939,17 @@ def disconnect_data_item(data_item_: typing.Optional[Persistence.PersistentObjec
             connect_data_item(typing.cast(DataItem.DataItem, self.__data_item_reference.item))
 
     def close(self) -> None:
-        # wait for display values threads to finish. first notify the thread that we are closing, then wait for it
-        # to complete by getting the future and waiting for it to complete. then clear the streams to release any
-        # resources (display values).
+        # ensure display data channel is disconnected from the data item in case it was never added to document
+        # and subsequently about_to_be_removed is not called.
+        if self.__is_data_item_connected and (data_item := self.data_item):
+            data_item.remove_display_data_channel(self)
+            self.__is_data_item_connected = False
         self.__data_item_reference.on_item_registered = None
         self.__data_item_reference.on_item_unregistered = None
         self.__closing = True
+        # wait for display values threads to finish. first notify the thread that we are closing, then wait for it
+        # to complete by getting the future and waiting for it to complete. then clear the streams to release any
+        # resources (display values).
         with self.__display_values_update_lock:
             display_values_future = self.__display_values_future
         if display_values_future:
@@ -960,6 +969,7 @@ def about_to_be_removed(self, container: Persistence.PersistentObject) -> None:
         # tell the data item that this display data channel is no longer referencing it
         if self.__data_item:
             self.__data_item.remove_display_data_channel(self)
+            self.__is_data_item_connected = False
         self.notify_property_changed("display_item")  # used for implicit connections
         super().about_to_be_removed(container)
 

nion/swift/test/DisplayItem_test.py

@@ -53,6 +53,17 @@ def test_display_item_with_multiple_display_data_channels_has_sensible_propertie
             self.assertIsNotNone(display_item.project_str)
             self.assertIsNotNone(display_item.used_display_type)
 
+    def test_display_data_channel_disconnects_if_display_item_closed(self):
+        with TestContext.create_memory_context() as test_context:
+            document_model = test_context.create_document_model()
+            data_item = DataItem.DataItem(numpy.zeros((8, 8), numpy.uint32))
+            document_model.append_data_item(data_item)
+            display_item = document_model.get_display_item_for_data_item(data_item)
+            display_item_snapshot = display_item.snapshot()
+            self.assertEqual(2, len(data_item.display_data_channels))
+            display_item_snapshot.close()
+            self.assertEqual(1, len(data_item.display_data_channels))
+
     def test_display_item_snapshot_and_copy_preserve_display_type(self):
         with TestContext.create_memory_context() as test_context:
             document_model = test_context.create_document_model()