feat: Add InMemoryTarget to store data in a memory instead of dumping to files (#432)

* feat: add in-memory target

* feat: add a data format to be stored in `Repository`
feat: last_modification_time feature in `InMemoryTarget`
style: add some type hints
fix: fix typo in `InMemoryCacheRepository`
test: add some tests for `InMemoryTarget` and `InMemoryCacheRepository`

* fix: fix linting errors

* fix: update type union shorthand to  to make compatible with py39

* style: refactor some base classes to inherite from

* fix: remove unnessesary optional type

* fix: fix format error

* chore: add an assertion error message
style: update variable name from `id` to `key`

* style: update the variable name  to  for code consistency

* docs: add a document of how to create InMemoryTarget

* chore: update name from `inmemory` to `in_memory`
chore: add type hints
style: remove `Protocol`

* fix: fix lint errors

* chore: raise Error when using InMemoryTarget with Redis

* fix: fix a linting error

---------

Co-authored-by: Ibuki Nakamura <[email protected]>

Author: utotsubasa

docs/task_on_kart.rst

@@ -286,3 +286,26 @@ If you want to dump csv file with other encodings, you can use `encoding` parame
     def output(self):
         return self.make_target('file_name.csv', processor=CsvFileProcessor(encoding='cp932'))
         # This will dump csv as 'cp932' which is used in Windows.
+
+Cache output in memory instead of dumping to files
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+You can use :class:`~InMemoryTarget` to cache output in memory instead of dumping to files by calling :func:`~gokart.target.make_in_memory_target`.
+
+Please note that :class:`~InMemoryTarget` is an experimental feature.
+
+.. code:: python
+
+    from gokart.in_memory.target import make_in_memory_target
+
+    def output(self):
+        unique_id = self.make_unique_id() if use_unique_id else None
+        # TaskLock is not supported in InMemoryTarget, so it's dummy
+        task_lock_params = make_task_lock_params(
+            file_path='dummy_path',
+            unique_id=unique_id,
+            redis_host=None,
+            redis_port=None,
+            redis_timeout=self.redis_timeout,
+            raise_task_lock_exception_on_collision=False,
+        )
+        return make_in_memory_target('dummy_path', task_lock_params, unique_id)

gokart/in_memory/__init__.py

@@ -0,0 +1,2 @@
+from .repository import InMemoryCacheRepository  # noqa:F401
+from .target import InMemoryTarget, make_in_memory_target  # noqa:F401

gokart/in_memory/data.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+
+
+@dataclass
+class InMemoryData:
+    value: Any
+    last_modification_time: datetime
+
+    @classmethod
+    def create_data(self, value: Any) -> 'InMemoryData':
+        return InMemoryData(value=value, last_modification_time=datetime.now())

gokart/in_memory/repository.py

@@ -0,0 +1,44 @@
+from typing import Any, Iterator
+
+from .data import InMemoryData
+
+
+class InMemoryCacheRepository:
+    _cache: dict[str, InMemoryData] = {}
+
+    def __init__(self):
+        pass
+
+    def get_value(self, key: str) -> Any:
+        return self._get_data(key).value
+
+    def get_last_modification_time(self, key: str):
+        return self._get_data(key).last_modification_time
+
+    def _get_data(self, key: str) -> InMemoryData:
+        return self._cache[key]
+
+    def set_value(self, key: str, obj: Any) -> None:
+        data = InMemoryData.create_data(obj)
+        self._cache[key] = data
+
+    def has(self, key: str) -> bool:
+        return key in self._cache
+
+    def remove(self, key: str) -> None:
+        assert self.has(key), f'{key} does not exist.'
+        del self._cache[key]
+
+    def empty(self) -> bool:
+        return not self._cache
+
+    def clear(self) -> None:
+        self._cache.clear()
+
+    def get_gen(self) -> Iterator[tuple[str, Any]]:
+        for key, data in self._cache.items():
+            yield key, data.value
+
+    @property
+    def size(self) -> int:
+        return len(self._cache)

gokart/in_memory/target.py

@@ -0,0 +1,45 @@
+from datetime import datetime
+from typing import Any
+
+from gokart.in_memory.repository import InMemoryCacheRepository
+from gokart.target import TargetOnKart, TaskLockParams
+
+_repository = InMemoryCacheRepository()
+
+
+class InMemoryTarget(TargetOnKart):
+    def __init__(self, data_key: str, task_lock_param: TaskLockParams):
+        if task_lock_param.should_task_lock:
+            raise ValueError('Redis with `InMemoryTarget` is not currently supported.')
+
+        self._data_key = data_key
+        self._task_lock_params = task_lock_param
+
+    def _exists(self) -> bool:
+        return _repository.has(self._data_key)
+
+    def _get_task_lock_params(self) -> TaskLockParams:
+        return self._task_lock_params
+
+    def _load(self) -> Any:
+        return _repository.get_value(self._data_key)
+
+    def _dump(self, obj: Any) -> None:
+        return _repository.set_value(self._data_key, obj)
+
+    def _remove(self) -> None:
+        _repository.remove(self._data_key)
+
+    def _last_modification_time(self) -> datetime:
+        if not _repository.has(self._data_key):
+            raise ValueError(f'No object(s) which id is {self._data_key} are stored before.')
+        time = _repository.get_last_modification_time(self._data_key)
+        return time
+
+    def _path(self) -> str:
+        # TODO: this module name `_path` migit not be appropriate
+        return self._data_key
+
+
+def make_in_memory_target(target_key: str, task_lock_params: TaskLockParams) -> InMemoryTarget:
+    return InMemoryTarget(target_key, task_lock_params)

test/in_memory/test_in_memory_target.py

@@ -0,0 +1,56 @@
+from datetime import datetime
+from time import sleep
+
+import pytest
+
+from gokart.conflict_prevention_lock.task_lock import TaskLockParams
+from gokart.in_memory import InMemoryCacheRepository, InMemoryTarget, make_in_memory_target
+
+
+class TestInMemoryTarget:
+    @pytest.fixture
+    def task_lock_params(self) -> TaskLockParams:
+        return TaskLockParams(
+            redis_host=None,
+            redis_port=None,
+            redis_timeout=None,
+            redis_key='dummy',
+            should_task_lock=False,
+            raise_task_lock_exception_on_collision=False,
+            lock_extend_seconds=0,
+        )
+
+    @pytest.fixture
+    def target(self, task_lock_params: TaskLockParams) -> InMemoryTarget:
+        return make_in_memory_target(target_key='dummy_key', task_lock_params=task_lock_params)
+
+    @pytest.fixture(autouse=True)
+    def clear_repo(self) -> None:
+        InMemoryCacheRepository().clear()
+
+    def test_dump_and_load_data(self, target: InMemoryTarget):
+        dumped = 'dummy_data'
+        target.dump(dumped)
+        loaded = target.load()
+        assert loaded == dumped
+
+    def test_exist(self, target: InMemoryTarget):
+        assert not target.exists()
+        target.dump('dummy_data')
+        assert target.exists()
+
+    def test_last_modified_time(self, target: InMemoryTarget):
+        input = 'dummy_data'
+        target.dump(input)
+        time = target.last_modification_time()
+        assert isinstance(time, datetime)
+
+        sleep(0.1)
+        another_input = 'another_data'
+        target.dump(another_input)
+        another_time = target.last_modification_time()
+        assert time < another_time
+
+        target.remove()
+        with pytest.raises(ValueError):
+            assert target.last_modification_time()

test/in_memory/test_repository.py

@@ -0,0 +1,63 @@
+import time
+
+import pytest
+
+from gokart.in_memory import InMemoryCacheRepository as Repo
+
+dummy_num = 100
+
+
+class TestInMemoryCacheRepository:
+    @pytest.fixture
+    def repo(self) -> Repo:
+        repo = Repo()
+        repo.clear()
+        return repo
+
+    def test_set(self, repo: Repo):
+        repo.set_value('dummy_key', dummy_num)
+        assert repo.size == 1
+        for key, value in repo.get_gen():
+            assert (key, value) == ('dummy_key', dummy_num)
+
+        repo.set_value('another_key', 'another_value')
+        assert repo.size == 2
+
+    def test_get(self, repo: Repo):
+        repo.set_value('dummy_key', dummy_num)
+        repo.set_value('another_key', 'another_value')
+
+        """Raise Error when key doesn't exist."""
+        with pytest.raises(KeyError):
+            repo.get_value('not_exist_key')
+
+        assert repo.get_value('dummy_key') == dummy_num
+        assert repo.get_value('another_key') == 'another_value'
+
+    def test_empty(self, repo: Repo):
+        assert repo.empty()
+        repo.set_value('dummmy_key', dummy_num)
+        assert not repo.empty()
+
+    def test_has(self, repo: Repo):
+        assert not repo.has('dummy_key')
+        repo.set_value('dummy_key', dummy_num)
+        assert repo.has('dummy_key')
+        assert not repo.has('not_exist_key')
+
+    def test_remove(self, repo: Repo):
+        repo.set_value('dummy_key', dummy_num)
+
+        with pytest.raises(AssertionError):
+            repo.remove('not_exist_key')
+
+        repo.remove('dummy_key')
+        assert not repo.has('dummy_key')
+
+    def test_last_modification_time(self, repo: Repo):
+        repo.set_value('dummy_key', dummy_num)
+        date1 = repo.get_last_modification_time('dummy_key')
+        time.sleep(0.1)
+        repo.set_value('dummy_key', dummy_num)
+        date2 = repo.get_last_modification_time('dummy_key')
+        assert date1 < date2