🐛 FIX: aiida/repository typing (#4920)

There was a subtle bug in `_insert_file`, where in the else clause of
the conditional `self.get_directory` was missing the parentheses to call
the method. This didn't show up in the tests because the code path was
actually not reachable since the conditional is always True. The original
intention was to get the parent directory of the file to be created, and
only create it if it wasn't already created. However, this behavior is
anyway guaranteed by `create_directory` so this can be safely called in
all instances and we can get rid of the entire conditional.

In certain methods, such as `open`, an explicit assert had to be placed
to check that the `File` instance returned by `get_file` had a `key`
that is not `None`, or `mypy` would complain. The reason that this is an
assert is because a `File` of type `FileType.FILE` always has `key` that
is not `None`, as is guaranteed by its constructor and the `get_file`
method will check that the path does correspond to a file and not a
directory, in which case it would raise `IsADirectory`. However, `mypy`
cannot introspect this polymorphism and so the explicit check is needed.
Since this would constitute an internal coding error if the check fails,
we use an assert.

Finally, the design of the `put_object_from_filelike` method in the
abstract class that performed the validation such that subclasses would
not have to, also didn't please the `mypy` gods. So instead, the
validation implementation of this method is moved to the protected method
`_put_object_from_filelike` and the base class performs the validation
and then calls this method.

.pre-commit-config.yaml

@@ -48,6 +48,7 @@ repos:
                 aiida/manage/database/delete/nodes.py|
                 aiida/orm/nodes/node.py|
                 aiida/orm/nodes/process/.*py|
+                aiida/repository/.*py|
                 aiida/tools/graph/graph_traversers.py|
                 aiida/tools/groups/paths.py|
                 aiida/tools/importexport/archive/.*py|

aiida/backends/general/migrations/utils.py

@@ -108,7 +108,7 @@ def is_initialised(self) -> bool:
     def erase(self):
         raise NotImplementedError()
 
-    def put_object_from_filelike(self, handle: io.BufferedIOBase) -> str:
+    def _put_object_from_filelike(self, handle: io.BufferedIOBase) -> str:
         """Store the byte contents of a file in the repository.
 
         :param handle: filelike object with the byte content to be stored.

aiida/orm/nodes/repository.py

@@ -105,7 +105,7 @@ def list_object_names(self, path: str = None) -> typing.List[str]:
         return self._repository.list_object_names(path)
 
     @contextlib.contextmanager
-    def open(self, path: str, mode='r') -> io.BufferedReader:
+    def open(self, path: str, mode='r') -> typing.Iterator[typing.BinaryIO]:
         """Open a file handle to an object stored under the given key.
 
         .. note:: this should only be used to open a handle to read an existing file. To write a new file use the method

aiida/repository/__init__.py

@@ -13,4 +13,4 @@
 from .common import *
 from .repository import *
 
-__all__ = (backend.__all__ + common.__all__ + repository.__all__)
+__all__ = (backend.__all__ + common.__all__ + repository.__all__)  # type: ignore[name-defined]

aiida/repository/backend/__init__.py

@@ -5,4 +5,4 @@
 from .disk_object_store import *
 from .sandbox import *
 
-__all__ = (abstract.__all__ + disk_object_store.__all__ + sandbox.__all__)
+__all__ = (abstract.__all__ + disk_object_store.__all__ + sandbox.__all__)  # type: ignore[name-defined]

aiida/repository/backend/abstract.py

@@ -44,7 +44,7 @@ def is_initialised(self) -> bool:
         """Return whether the repository has been initialised."""
 
     @abc.abstractmethod
-    def erase(self):
+    def erase(self) -> None:
         """Delete the repository itself and all its contents.
 
         .. note:: This should not merely delete the contents of the repository but any resources it created. For
@@ -53,17 +53,23 @@ def erase(self):
         """
 
     @staticmethod
-    def is_readable_byte_stream(handle):
+    def is_readable_byte_stream(handle) -> bool:
         return hasattr(handle, 'read') and hasattr(handle, 'mode') and 'b' in handle.mode
 
-    def put_object_from_filelike(self, handle: io.BufferedIOBase) -> str:
+    def put_object_from_filelike(self, handle: typing.BinaryIO) -> str:
         """Store the byte contents of a file in the repository.
 
         :param handle: filelike object with the byte content to be stored.
         :return: the generated fully qualified identifier for the object within the repository.
+        :raises TypeError: if the handle is not a byte stream.
         """
         if not isinstance(handle, io.BytesIO) and not self.is_readable_byte_stream(handle):
             raise TypeError(f'handle does not seem to be a byte stream: {type(handle)}.')
+        return self._put_object_from_filelike(handle)
+
+    @abc.abstractmethod
+    def _put_object_from_filelike(self, handle: typing.BinaryIO) -> str:
+        pass
 
     def put_object_from_file(self, filepath: typing.Union[str, pathlib.Path]) -> str:
         """Store a new object with contents of the file located at `filepath` on this file system.
@@ -84,7 +90,7 @@ def has_object(self, key: str) -> bool:
         """
 
     @contextlib.contextmanager
-    def open(self, key: str) -> io.BufferedIOBase:
+    def open(self, key: str) -> typing.Iterator[typing.BinaryIO]:  # type: ignore[return]
         """Open a file handle to an object stored under the given key.
 
         .. note:: this should only be used to open a handle to read an existing file. To write a new file use the method

aiida/repository/backend/disk_object_store.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 """Implementation of the ``AbstractRepositoryBackend`` using the ``disk-objectstore`` as the backend."""
 import contextlib
-import io
 import shutil
 import typing
 
@@ -23,7 +22,9 @@ def __init__(self, container):
 
     def __str__(self) -> str:
         """Return the string representation of this repository."""
-        return f'DiskObjectStoreRepository: {self.container.container_id} | {self.container.get_folder()}'
+        if self.is_initialised:
+            return f'DiskObjectStoreRepository: {self.container.container_id} | {self.container.get_folder()}'
+        return 'DiskObjectStoreRepository: <uninitialised>'
 
     @property
     def uuid(self) -> typing.Optional[str]:
@@ -45,7 +46,7 @@ def is_initialised(self) -> bool:
         return self.container.is_initialised
 
     @property
-    def container(self):
+    def container(self) -> Container:
         return self._container
 
     def erase(self):
@@ -55,14 +56,13 @@ def erase(self):
         except FileNotFoundError:
             pass
 
-    def put_object_from_filelike(self, handle: io.BufferedIOBase) -> str:
+    def _put_object_from_filelike(self, handle: typing.BinaryIO) -> str:
         """Store the byte contents of a file in the repository.
 
         :param handle: filelike object with the byte content to be stored.
         :return: the generated fully qualified identifier for the object within the repository.
         :raises TypeError: if the handle is not a byte stream.
         """
-        super().put_object_from_filelike(handle)
         return self.container.add_object(handle.read())
 
     def has_object(self, key: str) -> bool:
@@ -74,7 +74,7 @@ def has_object(self, key: str) -> bool:
         return self.container.has_object(key)
 
     @contextlib.contextmanager
-    def open(self, key: str) -> io.BufferedIOBase:
+    def open(self, key: str) -> typing.Iterator[typing.BinaryIO]:
         """Open a file handle to an object stored under the given key.
 
         .. note:: this should only be used to open a handle to read an existing file. To write a new file use the method

aiida/repository/backend/sandbox.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 """Implementation of the ``AbstractRepositoryBackend`` using a sandbox folder on disk as the backend."""
 import contextlib
-import io
 import os
 import shutil
 import typing
@@ -16,11 +15,14 @@ class SandboxRepositoryBackend(AbstractRepositoryBackend):
     """Implementation of the ``AbstractRepositoryBackend`` using a sandbox folder on disk as the backend."""
 
     def __init__(self):
-        self._sandbox = None
+        from aiida.common.folders import SandboxFolder
+        self._sandbox: typing.Optional[SandboxFolder] = None
 
     def __str__(self) -> str:
         """Return the string representation of this repository."""
-        return f'SandboxRepository: {self._sandbox.abspath}'
+        if self.is_initialised:
+            return f'SandboxRepository: {self._sandbox.abspath if self._sandbox else "null"}'
+        return 'SandboxRepository: <uninitialised>'
 
     def __del__(self):
         """Delete the entire sandbox folder if it was instantiated and still exists."""
@@ -68,15 +70,13 @@ def erase(self):
             finally:
                 self._sandbox = None
 
-    def put_object_from_filelike(self, handle: io.BufferedIOBase) -> str:
+    def _put_object_from_filelike(self, handle: typing.BinaryIO) -> str:
         """Store the byte contents of a file in the repository.
 
         :param handle: filelike object with the byte content to be stored.
         :return: the generated fully qualified identifier for the object within the repository.
         :raises TypeError: if the handle is not a byte stream.
         """
-        super().put_object_from_filelike(handle)
-
         key = str(uuid.uuid4())
         filepath = os.path.join(self.sandbox.abspath, key)
 
@@ -94,7 +94,7 @@ def has_object(self, key: str) -> bool:
         return key in os.listdir(self.sandbox.abspath)
 
     @contextlib.contextmanager
-    def open(self, key: str) -> io.BufferedIOBase:
+    def open(self, key: str) -> typing.Iterator[typing.BinaryIO]:
         """Open a file handle to an object stored under the given key.
 
         .. note:: this should only be used to open a handle to read an existing file. To write a new file use the method