def test_remove():
"""Unit test for method remove."""
index_to_remove = 8
temp_file = FileIO.safe_temp_file(file_name='pupyl.index')
temp_dir = os.path.dirname(temp_file)
with Index(TEST_VECTOR_SIZE, data_dir=temp_dir) as index:
for _ in range(16):
index.append(numpy.random.normal(size=TEST_VECTOR_SIZE))
test_size_before = len(index)
test_value = index[index_to_remove]
with Index(TEST_VECTOR_SIZE, data_dir=temp_dir) as index:
index.remove(index_to_remove)
assert len(index) == test_size_before - 1
numpy.testing.assert_raises(
AssertionError,
numpy.testing.assert_array_equal,
test_value,
index[index_to_remove]
)
def __init__(self, size, data_dir=None, trees=.001, volatile=False):
"""
Indexing tensors operations and nearest neighbours search.
Parameters
----------
size: int
Shape of unidimensional vectors which will be indexed
data_dir: str
Location where to load or save the index
trees (optional): float
Defines the number of trees to create based on the dataset
size. Should be a number between 0 and 1.
volatile (optional): bool
If the index will be temporary or not.
"""
self._position = -1
self._size = size
self._data_dir = data_dir
self._trees = trees
self._volatile = volatile
if self._data_dir and not self._volatile:
if os.path.isfile(self._data_dir):
raise OSError('data_dir parameter is not a directory')
os.makedirs(self._data_dir, exist_ok=True)
self._path = os.path.join(self._data_dir, self.index_name)
elif not self._data_dir and not self._volatile:
raise NoDataDirForPermanentIndex
elif not self._data_dir and self._volatile:
_temp_file = FileIO.safe_temp_file()
self._data_dir = os.path.dirname(_temp_file)
self._path = _temp_file
else:
raise DataDirDefinedForVolatileIndex
if os.path.isfile(self._path):
try:
self.tree = AnnoyIndex(size, metric='angular')
self.tree.load(self._path)
self._is_new_index = False
except OSError as os_error:
raise FileIsNotAnIndex from os_error
else:
self.tree = AnnoyIndex(size, metric='angular')
self._is_new_index = True
self._image_database = ImageDatabase(
import_images=True,
data_dir=self._data_dir,
)
def test_safe_temp_file_exists():
"""Unit test for method safe_temp_file, file exists case."""
test_temp_file_name = 'just_a_temp_file.txt'
Path(join(tempfile.gettempdir(), test_temp_file_name)).touch()
_ = FileIO.safe_temp_file(file_name=test_temp_file_name)
assert not exists(test_temp_file_name)
def test_pop():
"""Unit test for method pop."""
temp_file = FileIO.safe_temp_file(file_name='pupyl.index')
temp_dir = os.path.dirname(temp_file)
with Index(TEST_VECTOR_SIZE, data_dir=temp_dir) as index:
for _ in range(16):
index.append(numpy.random.normal(size=TEST_VECTOR_SIZE))
test_size_before = len(index)
test_value_before = index[-1]
with Index(TEST_VECTOR_SIZE, data_dir=temp_dir) as index:
test_value_after = index.pop()
assert len(index) == test_size_before - 1
numpy.testing.assert_array_equal(
test_value_before,
test_value_after
)
def test_safe_temp_file():
"""Unit test for method safe_temp_file."""
test_temp_file_name = FileIO.safe_temp_file()
assert not exists(test_temp_file_name)