diff --git a/docarray/array/base.py b/docarray/array/base.py index d136eb1e0b7..2f333deffb6 100644 --- 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 index ab49d0f4e22..1a2c9866662 100644 --- 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 index 3ca78acc0cc..9c2007b3c7b 100644 --- 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 index 458e563f7c5..e4e59c385ce 100644 --- 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 == ""