Implement trio.abc.AsyncResource (ie: aclose()) for all datastores

Lot's of “code change” unfortunately, since all the tests were convert to
properly utilize this facility everywhere, which ended indenting just about
every test line by +1.

Author: ntninja

datastore/adapter/_support.py

@@ -1,5 +1,7 @@
 import typing
 
+import trio
+
 import datastore
 
 T_co = typing.TypeVar("T_co", covariant=True)
@@ -58,3 +60,24 @@ def remove_datastore(self, store: DS) -> None:
 	def insert_datastore(self, index: int, store: DS) -> None:
 		"""Inserts datastore `store` into this collection at `index`."""
 		self._stores.insert(index, store)
+	
+	async def _stores_cleanup(self) -> None:
+		"""Closes and removes all added datastores"""
+		errors: typing.List[Exception] = []
+		
+		while len(self._stores):
+			store = self._stores.pop()
+			
+			try:
+				await store.aclose()
+			except trio.Cancelled:
+				pass  # We check for cancellation later on
+			except Exception as error:
+				errors.append(error)
+		
+		# Ensure error propagation
+		if errors:
+			raise trio.MultiError(errors)
+		
+		# Ensure cancellation is propagated
+		await trio.sleep(0)

datastore/adapter/sharded.py

@@ -92,6 +92,11 @@ def _shard_query_generator(self, query):
 
 				if shard_query.limit <= 0:
 					break  # we're already done!
+	
+	
+	async def aclose(self) -> None:
+		"""Closes and removes all added datastores"""
+		await self._stores_cleanup()
 
 
 

datastore/adapter/tiered.py

@@ -143,6 +143,11 @@ async def contains(self, key: datastore.Key) -> bool:
 			if await store.contains(key):
 				return True
 		return False
+	
+	
+	async def aclose(self) -> None:
+		"""Closes and removes all added datastores"""
+		await self._stores_cleanup()
 
 
 class BinaryAdapter(

datastore/core/binarystore.py

@@ -29,7 +29,7 @@ def is_valid_value_type(value: util.stream.ArbitraryReceiveStream) -> bool:
 
 
 
-class Datastore:
+class Datastore(trio.abc.AsyncResource):
 	"""A Datastore represents storage for any string key to binary value pair.
 
 	Datastores are general enough to be backed by all kinds of different storage:
@@ -166,6 +166,16 @@ async def get_all(self, key: key_.Key) -> bytes:
 			An internal error occurred
 		"""
 		return await (await self.get(key)).collect()
+	
+	
+	async def aclose(self) -> None:
+		"""Closes this any resources held by this datastore, possibly blocking
+		
+		Carefully read the documentation of :class:`trio.abc.AsyncResource`,
+		particularily with regards to concellation and forceful closings, when
+		implementating this.
+		"""
+		pass
 
 
 
@@ -284,6 +294,12 @@ async def contains(self, key: key_.Key) -> bool:
 
 	def __len__(self) -> int:
 		return sum(map(len, self._items.values()))
+	
+	
+	async def aclose(self) -> None:
+		"""Deletes all items from this datastore"""
+		self._items.clear()
+		await super().aclose()
 
 
 
@@ -406,6 +422,19 @@ async def contains(self, key: key_.Key) -> bool:
 			return await self.child_datastore.contains(key)
 		else:
 			return await Datastore.contains(self, key)
+	
+	
+	async def aclose(self) -> None:
+		"""Closes this any resources held by the child datastore
+		
+		Carefully read the documentation of :class:`trio.abc.AsyncResource`,
+		particularily with regards to concellation and forceful closings, when
+		implementating this.
+		"""
+		try:
+			await self.child_datastore.aclose()
+		finally:
+			await super().aclose()
 
 
 Datastore.ADAPTER_T = Adapter

datastore/core/objectstore.py

@@ -32,7 +32,7 @@ def is_valid_value_type(value: util.stream.ArbitraryReceiveChannel) -> bool:
 	)) and not isinstance(value, (str, bytes))
 
 
-class Datastore(typing.Generic[T_co]):
+class Datastore(typing.Generic[T_co], trio.abc.AsyncResource):
 	"""A Datastore represents storage for any string key to an arbitrary Python object pair.
 
 	Datastores are general enough to be backed by all kinds of different storage:
@@ -193,6 +193,16 @@ async def get_all(self, key: key_.Key) -> typing.List[T_co]:
 			An internal error occurred
 		"""
 		return await (await self.get(key)).collect()
+	
+	
+	async def aclose(self) -> None:
+		"""Closes this any resources held by this datastore, possibly blocking
+		
+		Carefully read the documentation of :class:`trio.abc.AsyncResource`,
+		particularily with regards to concellation and forceful closings, when
+		implementating this.
+		"""
+		pass
 
 
 
@@ -336,6 +346,12 @@ async def query(self, query: query_.Query) -> query_.Cursor:
 
 	def __len__(self) -> int:
 		return sum(map(len, self._items.values()))
+	
+	
+	async def aclose(self) -> None:
+		"""Deletes all items from this datastore"""
+		self._items.clear()
+		await super().aclose()
 
 
 
@@ -483,6 +499,19 @@ async def contains(self, key: key_.Key) -> bool:
 			return await self.child_datastore.contains(key)
 		else:
 			return await Datastore.contains(self, key)
+	
+	
+	async def aclose(self) -> None:
+		"""Closes this any resources held by the child datastore
+		
+		Carefully read the documentation of :class:`trio.abc.AsyncResource`,
+		particularily with regards to concellation and forceful closings, when
+		implementating this.
+		"""
+		try:
+			await self.child_datastore.aclose()
+		finally:
+			await super().aclose()
 
 
 Datastore.ADAPTER_T = Adapter

tests/adapter/test_directory.py

@@ -1,3 +1,4 @@
+import contextlib
 import logging
 import typing
 
@@ -10,9 +11,14 @@
 
 @trio.testing.trio_test
 async def test_dir_simple(DatastoreTests):
-	s1 = datastore.adapter.directory.ObjectDatastore(datastore.ObjectDictDatastore())
-	s2 = datastore.adapter.directory.ObjectDatastore(datastore.ObjectDictDatastore())
-	await DatastoreTests([s1, s2]).subtest_simple()
+	async with contextlib.AsyncExitStack() as stack:
+		s1 = stack.push_async_exit(
+			datastore.adapter.directory.ObjectDatastore(datastore.ObjectDictDatastore())
+		)
+		s2 = stack.push_async_exit(
+			datastore.adapter.directory.ObjectDatastore(datastore.ObjectDictDatastore())
+		)
+		await DatastoreTests([s1, s2]).subtest_simple()
 
 
 ##########################
@@ -29,106 +35,103 @@ class ObjectDirectoryDictDatastore(
 
 @trio.testing.trio_test
 async def test_directory_simple(DatastoreTests):
-	s1 = ObjectDirectoryDictDatastore()
-	s2 = ObjectDirectoryDictDatastore()
-	await DatastoreTests([s1, s2]).subtest_simple()
+	async with contextlib.AsyncExitStack() as stack:
+		s1 = stack.push_async_exit(ObjectDirectoryDictDatastore())
+		s2 = stack.push_async_exit(ObjectDirectoryDictDatastore())
+		await DatastoreTests([s1, s2]).subtest_simple()
 
 
 @trio.testing.trio_test
 async def test_directory_init():
-	ds = ObjectDirectoryDictDatastore()
+	async with ObjectDirectoryDictDatastore() as ds:
+		# initialize directory at /foo
+		dir_key = datastore.Key('/foo')
+		await ds.directory(dir_key)
+		assert await ds.get_all(dir_key) == []
 
-	# initialize directory at /foo
-	dir_key = datastore.Key('/foo')
-	await ds.directory(dir_key)
-	assert await ds.get_all(dir_key) == []
+		# can add to dir
+		bar_key = datastore.Key('/foo/bar')
+		await ds.directory_add(dir_key, bar_key)
+		assert await ds.get_all(dir_key) == [str(bar_key)]
 
-	# can add to dir
-	bar_key = datastore.Key('/foo/bar')
-	await ds.directory_add(dir_key, bar_key)
-	assert await ds.get_all(dir_key) == [str(bar_key)]
-
-	# re-init does not wipe out directory at /foo
-	dir_key = datastore.Key('/foo')
-	with pytest.raises(KeyError):
-		await ds.directory(dir_key, exist_ok=False)
-	await ds.directory(dir_key, exist_ok=True)
-	assert await ds.get_all(dir_key) == [str(bar_key)]
+		# re-init does not wipe out directory at /foo
+		dir_key = datastore.Key('/foo')
+		with pytest.raises(KeyError):
+			await ds.directory(dir_key, exist_ok=False)
+		await ds.directory(dir_key, exist_ok=True)
+		assert await ds.get_all(dir_key) == [str(bar_key)]
 
 
 @trio.testing.trio_test
 async def test_directory_basic():
-	ds = ObjectDirectoryDictDatastore()
-
-	# initialize directory at /foo
-	dir_key = datastore.Key('/foo')
-	await ds.directory(dir_key)
-
-	# adding directory entries
-	bar_key = datastore.Key('/foo/bar')
-	baz_key = datastore.Key('/foo/baz')
-	await ds.directory_add(dir_key, bar_key)
-	await ds.directory_add(dir_key, baz_key)
-	keys = [key async for key in ds.directory_read(dir_key)]
-	assert keys == [bar_key, baz_key]
-
-	# removing directory entries
-	await ds.directory_remove(dir_key, bar_key)
-	keys = [key async for key in ds.directory_read(dir_key)]
-	assert keys == [baz_key]
-
-	await ds.directory_remove(dir_key, baz_key)
-	keys = [key async for key in ds.directory_read(dir_key)]
-	assert keys == []
-
-	# generator
-	with pytest.raises(StopAsyncIteration):
-		gen = ds.directory_read(dir_key).__aiter__()
-		await gen.__anext__()
+	async with ObjectDirectoryDictDatastore() as ds:
+		# initialize directory at /foo
+		dir_key = datastore.Key('/foo')
+		await ds.directory(dir_key)
+
+		# adding directory entries
+		bar_key = datastore.Key('/foo/bar')
+		baz_key = datastore.Key('/foo/baz')
+		await ds.directory_add(dir_key, bar_key)
+		await ds.directory_add(dir_key, baz_key)
+		keys = [key async for key in ds.directory_read(dir_key)]
+		assert keys == [bar_key, baz_key]
+
+		# removing directory entries
+		await ds.directory_remove(dir_key, bar_key)
+		keys = [key async for key in ds.directory_read(dir_key)]
+		assert keys == [baz_key]
+
+		await ds.directory_remove(dir_key, baz_key)
+		keys = [key async for key in ds.directory_read(dir_key)]
+		assert keys == []
+
+		# generator
+		with pytest.raises(StopAsyncIteration):
+			gen = ds.directory_read(dir_key).__aiter__()
+			await gen.__anext__()
 
 
 @trio.testing.trio_test
 async def test_directory_double_add():
-	ds = ObjectDirectoryDictDatastore()
-
-	# initialize directory at /foo
-	dir_key = datastore.Key('/foo')
-	await ds.directory(dir_key)
+	async with ObjectDirectoryDictDatastore() as ds:
+		# initialize directory at /foo
+		dir_key = datastore.Key('/foo')
+		await ds.directory(dir_key)
 
-	# adding directory entries
-	bar_key = datastore.Key('/foo/bar')
-	baz_key = datastore.Key('/foo/baz')
-	await ds.directory_add(dir_key, bar_key)
-	await ds.directory_add(dir_key, baz_key)
-	await ds.directory_add(dir_key, bar_key)
-	await ds.directory_add(dir_key, baz_key)
-	await ds.directory_add(dir_key, baz_key)
-	await ds.directory_add(dir_key, bar_key)
+		# adding directory entries
+		bar_key = datastore.Key('/foo/bar')
+		baz_key = datastore.Key('/foo/baz')
+		await ds.directory_add(dir_key, bar_key)
+		await ds.directory_add(dir_key, baz_key)
+		await ds.directory_add(dir_key, bar_key)
+		await ds.directory_add(dir_key, baz_key)
+		await ds.directory_add(dir_key, baz_key)
+		await ds.directory_add(dir_key, bar_key)
 
-	keys = [key async for key in ds.directory_read(dir_key)]
-	assert keys == [bar_key, baz_key]
+		keys = [key async for key in ds.directory_read(dir_key)]
+		assert keys == [bar_key, baz_key]
 
 
 @trio.testing.trio_test
 async def test_directory_remove():
-	ds = ObjectDirectoryDictDatastore()
-
-	# initialize directory at /foo
-	dir_key = datastore.Key('/foo')
-	await ds.directory(dir_key)
-
-	# adding directory entries
-	bar_key = datastore.Key('/foo/bar')
-	baz_key = datastore.Key('/foo/baz')
-	await ds.directory_add(dir_key, bar_key)
-	await ds.directory_add(dir_key, baz_key)
-	keys = [key async for key in ds.directory_read(dir_key)]
-	assert keys == [bar_key, baz_key]
-
-	# removing directory entries
-	await ds.directory_remove(dir_key, bar_key)
-	await ds.directory_remove(dir_key, bar_key, missing_ok=True)
-	with pytest.raises(KeyError):
-		await ds.directory_remove(dir_key, bar_key, missing_ok=False)
-	keys = [key async for key in ds.directory_read(dir_key)]
-	assert keys == [baz_key]
+	async with ObjectDirectoryDictDatastore() as ds:
+		# initialize directory at /foo
+		dir_key = datastore.Key('/foo')
+		await ds.directory(dir_key)
+
+		# adding directory entries
+		bar_key = datastore.Key('/foo/bar')
+		baz_key = datastore.Key('/foo/baz')
+		await ds.directory_add(dir_key, bar_key)
+		await ds.directory_add(dir_key, baz_key)
+		keys = [key async for key in ds.directory_read(dir_key)]
+		assert keys == [bar_key, baz_key]
+
+		# removing directory entries
+		await ds.directory_remove(dir_key, bar_key)
+		await ds.directory_remove(dir_key, bar_key, missing_ok=True)
+		with pytest.raises(KeyError):
+			await ds.directory_remove(dir_key, bar_key, missing_ok=False)
+		keys = [key async for key in ds.directory_read(dir_key)]
+		assert keys == [baz_key]