docarray · Jackmin801 · Oct 11, 2022 · Oct 11, 2022
diff --git a/docarray/array/base.py b/docarray/array/base.py
@@ -1,3 +1,5 @@
+import atexit
+from weakref import WeakMethod
 from typing import MutableSequence, TYPE_CHECKING, Union, Iterable
 
 from docarray import Document
@@ -11,6 +13,17 @@ def __init__(self, *args, storage: str = 'memory', **kwargs):
         super().__init__()
         self._init_storage(*args, **kwargs)
         self._init_subindices(*args, **kwargs)
+        self._register_exit_hook()
+
+    def _register_exit_hook(self):
+        weakref_del = WeakMethod(self.__del__)
+
+        def exit_hook():
+            if weakref_del() is not None:
+                weakref_del()()
+
+        atexit.register(exit_hook)
+        self._deleted = False
 
     def __add__(self: 'T', other: Union['Document', Iterable['Document']]) -> 'T':
         v = type(self)(self)

diff --git a/docarray/array/storage/base/getsetdel.py b/docarray/array/storage/base/getsetdel.py
@@ -326,5 +326,6 @@ def _save_offset2ids(self):
         ...
 
     def __del__(self):
-        if hasattr(self, '_offset2ids'):
+        if hasattr(self, '_offset2ids') and not self._deleted:
             self._save_offset2ids()
+            self._deleted = True
diff --git a/docarray/array/storage/elastic/seqlike.py b/docarray/array/storage/elastic/seqlike.py
@@ -45,13 +45,6 @@ def __contains__(self, x: Union[str, 'Document']):
         else:
             return False
 
-    def __del__(self):
-        """Delete this :class:`DocumentArrayElastic` object"""
-        self._save_offset2ids()
-
-        # if not self._persist:
-        #    self._offset2ids.clear()
-
     def __repr__(self):
         """Return the string representation of :class:`DocumentArrayElastic` object
         :return: string representation of this object

diff --git a/tests/unit/array/test_sequence.py b/tests/unit/array/test_sequence.py
@@ -1,3 +1,5 @@
+import gc
+import sys
 import tempfile
 import uuid
 
@@ -260,3 +262,37 @@ def test_set_and_append(index, storage, config):
         da.append(Document(id='new_new'))
 
     assert da[:, 'id'] == ['0', 'new', '2', '3', '4', 'new_new']
+
+
+@pytest.mark.parametrize(
+    'da_cls,config',
+    [
+        (DocumentArrayInMemory, lambda: None),
+        (DocumentArraySqlite, lambda: None),
+        # DocArrays backed by Weaviate in 0.17.0 has a refcount bug
+        # Remove this comment when the bug is fixed
+        # (DocumentArrayWeaviate, lambda: WeaviateConfig(n_dim=3)),
+        (DocumentArrayQdrant, lambda: QdrantConfig(n_dim=3)),
+        (DocumentArrayElastic, lambda: ElasticConfig(n_dim=3)),
+        (DocumentArrayRedis, lambda: RedisConfig(n_dim=3)),
+    ],
+)
+def test_deletion(da_cls, config, start_storage, capsys):
+    # Mock Deletion
+    class MockDA(da_cls):
+        def __del__(self):
+            assert not self._deleted
+            print('Goodbye')
+            self._deleted = True
+
+    da = MockDA(config=config())
+    assert not da._deleted
+    assert (
+        sys.getrefcount(da) == 2
+    ), f'{da_cls.__name__} initialized with {sys.getrefcount(da) - 1} references'
+
+    del da
+    gc.collect()
+    output = capsys.readouterr()
+    assert output.out == "Goodbye\n"
+    assert output.err == ""