Source code for swh.objstorage.backends.in_memory

# Copyright (C) 2017-2022  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information

from typing import Dict, Iterator

from typing_extensions import Literal

from swh.objstorage.exc import Error, ObjNotFoundError
from swh.objstorage.interface import CompositeObjId, ObjId
from swh.objstorage.objstorage import ObjStorage, compute_hash, objid_to_default_hex

[docs]class InMemoryObjStorage(ObjStorage): """In-Memory objstorage. Intended for test purposes. """ PRIMARY_HASH: Literal["sha1"] = "sha1" def __init__(self, **args): super().__init__() self.state: Dict[bytes, bytes] = {}
[docs] def check_config(self, *, check_write): return True
def _state_key(self, obj_id: ObjId) -> bytes: if isinstance(obj_id, dict): return obj_id[self.PRIMARY_HASH] else: return obj_id def __contains__(self, obj_id: ObjId) -> bool: return self._state_key(obj_id) in self.state def __iter__(self) -> Iterator[CompositeObjId]: for id_ in sorted(self.state): yield {self.PRIMARY_HASH: id_}
[docs] def add(self, content: bytes, obj_id: ObjId, check_presence: bool = True) -> None: if check_presence and obj_id in self: return self.state[self._state_key(obj_id)] = content
[docs] def get(self, obj_id: ObjId) -> bytes: if obj_id not in self: raise ObjNotFoundError(obj_id) return self.state[self._state_key(obj_id)]
[docs] def check(self, obj_id: ObjId) -> None: if obj_id not in self: raise ObjNotFoundError(obj_id) if compute_hash(self.state[self._state_key(obj_id)]) != self._state_key(obj_id): raise Error("Corrupt object %s" % objid_to_default_hex(obj_id))
[docs] def delete(self, obj_id: ObjId): super().delete(obj_id) # Check delete permission if obj_id not in self: raise ObjNotFoundError(obj_id) self.state.pop(self._state_key(obj_id)) return True