Add an explicit Cache API

Author: Mauko Quiroga

openfisca_core/data_storage.py

@@ -1,7 +1,7 @@
 import abc
 import os
 import shutil
-from typing import Any, Dict, KeysView, Optional
+from typing import Any, Dict, KeysView, Optional, Union
 
 import numpy
 
@@ -51,11 +51,11 @@ def delete(self, period: Optional[Period] = None) -> None:
         ...
 
     @abc.abstractmethod
-    def known_periods(self) -> KeysView[Period]:
+    def memory_usage(self) -> Dict[str, int]:
         ...
 
     @abc.abstractmethod
-    def memory_usage(self) -> Dict[str, int]:
+    def known_periods(self) -> KeysView[Period]:
         ...
 
 
@@ -74,7 +74,7 @@ def cast_period(self, period: Optional[Period], eternal: bool) -> Period:
 
 
 class MemoryStorage(StorageLike):
-    """Responsible for storing and retrieving values in memory."""
+    """Low-level class responsible for storing and retrieving values in memory."""
 
     def get(self, state: Dict[Period, Any], key: Period) -> Any:
         return state.get(key)
@@ -90,7 +90,6 @@ def delete_all(self, state: StateType) -> dict:
         state.clear()
         return state
 
-    # TODO: test
     def memory_usage(self, state: StateType) -> Dict[str, int]:
         if not state:
             return {
@@ -110,14 +109,17 @@ def memory_usage(self, state: StateType) -> Dict[str, int]:
 
 
 class DiskStorage(StorageLike):
-    """Responsible for storing and retrieving values on disk."""
+    """Low-level class responsible for storing and retrieving values on disk."""
     directory: str
     preserve: bool
 
     def __init__(self, directory: str, preserve: bool) -> None:
         self.directory = directory
         self.preserve = preserve
 
+        if not os.path.isdir(self.directory):
+            os.makedirs(self.directory, exist_ok = True)
+
     def get(self, state: Dict[Period, Any], key: Period) -> Any:
         file = state.get(key)
 
@@ -153,7 +155,6 @@ def delete_all(self, state: StateType) -> dict:
         state.clear()
         return state
 
-    # TODO: test
     def memory_usage(self, state: StateType) -> Dict[str, int]:
         if not state:
             return {
@@ -203,21 +204,26 @@ def __del__(self) -> None:
             shutil.rmtree(parent_dir)
 
 
-class InMemoryStorage(CachingLike, SupportsPeriodCasting):
+StorageType = Union[MemoryStorage, DiskStorage]
+
+
+class Cache(CachingLike, SupportsPeriodCasting):
     """
-    Low-level class responsible for storing and retrieving calculated vectors in memory.
+    Explicit Cache API responsible of:
 
-    TODO: separate concerns between the caching API and the storing API.
+    * keeping cache state
+    * reading from storages
+    * writing to storages
     """
 
     state: StateType
+    storage: StorageType
     is_eternal: bool
-    storage: MemoryStorage
 
-    def __init__(self, is_eternal: bool = False) -> None:
+    def __init__(self, storage: StorageType, is_eternal: bool = False) -> None:
         self.state = {}
+        self.storage = storage
         self.is_eternal = is_eternal
-        self.storage = MemoryStorage()
 
     def get(self, period: Period) -> Any:
         casted: Period = self.cast_period(period, self.is_eternal)
@@ -235,75 +241,20 @@ def delete(self, period: Optional[Period] = None) -> None:
         casted: Period = self.cast_period(period, self.is_eternal)
         self.state = self.storage.delete(self.state, casted)
 
-    # TODO: test
-    def known_periods(self) -> KeysView[Period]:
-        return self.state.keys()
-
-    # TODO: test
     def memory_usage(self) -> Dict[str, int]:
         return self.storage.memory_usage(self.state)
 
-    def get_known_periods(self) -> KeysView[Period]:
-        raise ValueError("TODO: add a deprecation warning")
-
-    # TODO: decide what to do with this.
-    def get_memory_usage(self) -> Dict[str, int]:
-        raise ValueError("TODO: add a deprecation warning")
-
-
-class OnDiskStorage(CachingLike, SupportsPeriodCasting):
-    """
-    Low-level class responsible for storing and retrieving calculated vectors on disk.
-
-    TODO: separate concerns between the caching API and the storing API.
-    """
-
-    state: StateType
-    is_eternal: bool
-    storage: DiskStorage
-
-    def __init__(
-            self,
-            storage_dir: str,
-            is_eternal: bool = False,
-            preserve_storage_dir: bool = False,
-            ) -> None:
-        self.state = {}
-        self.is_eternal = is_eternal
-        self.storage = DiskStorage(storage_dir, preserve_storage_dir)
-
-    def get(self, period: Period) -> Any:
-        casted: Period = self.cast_period(period, self.is_eternal)
-        return self.storage.get(self.state, casted)
-
-    def put(self, value: Any, period: Period) -> None:
-        casted: Period = self.cast_period(period, self.is_eternal)
-        self.state = self.storage.put(self.state, casted, value)
-
-    def delete(self, period: Optional[Period] = None) -> None:
-        if period is None:
-            self.state = self.storage.delete_all(self.state)
-            return
-
-        casted: Period = self.cast_period(period, self.is_eternal)
-        self.state = self.storage.delete(self.state, casted)
-
-    # TODO: test
     def known_periods(self) -> KeysView[Period]:
         return self.state.keys()
 
-    # TODO: test
-    def memory_usage(self) -> Dict[str, int]:
-        return self.storage.memory_usage(self.state)
-
+    # TODO : test
     def get_known_periods(self) -> KeysView[Period]:
         raise ValueError("TODO: add a deprecation warning")
 
+    # TODO : test
     def get_memory_usage(self) -> Dict[str, int]:
         raise ValueError("TODO: add a deprecation warning")
 
+    # TODO : test
     def restore(self, state: StateType) -> StateType:
         raise ValueError("TODO: add a deprecation warning")
-
-    def __del__(self) -> None:
-        raise ValueError("TODO: add a deprecation warning")

openfisca_core/holders.py

@@ -9,7 +9,7 @@
 
 from openfisca_core import periods
 from openfisca_core.commons import empty_clone
-from openfisca_core.data_storage import InMemoryStorage, OnDiskStorage
+from openfisca_core.data_storage import Cache, DiskStorage, MemoryStorage
 from openfisca_core.errors import PeriodMismatchError
 from openfisca_core.indexed_enums import Enum
 from openfisca_core.periods import MONTH, YEAR, ETERNITY
@@ -27,7 +27,7 @@ def __init__(self, variable, population):
         self.population = population
         self.variable = variable
         self.simulation = population.simulation
-        self._memory_storage = InMemoryStorage(is_eternal = (self.variable.definition_period == ETERNITY))
+        self._memory_storage = Cache(MemoryStorage(), is_eternal = (self.variable.definition_period == ETERNITY))
 
         # By default, do not activate on-disk storage, or variable dropping
         self._disk_storage = None
@@ -60,13 +60,13 @@ def create_disk_storage(self, directory = None, preserve = False):
         if directory is None:
             directory = self.simulation.data_storage_dir
         storage_dir = os.path.join(directory, self.variable.name)
+
         if not os.path.isdir(storage_dir):
             os.mkdir(storage_dir)
-        return OnDiskStorage(
-            storage_dir,
-            is_eternal = (self.variable.definition_period == ETERNITY),
-            preserve_storage_dir = preserve
-            )
+
+        storage = DiskStorage(storage_dir, preserve)
+
+        return Cache(storage, is_eternal = (self.variable.definition_period == ETERNITY))
 
     def delete_arrays(self, period = None):
         """

openfisca_core/tools/simulation_dumper.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 from openfisca_core.simulations import Simulation
-from openfisca_core.data_storage import OnDiskStorage
+from openfisca_core.data_storage import Cache, DiskStorage
 from openfisca_core.periods import ETERNITY
 
 
@@ -106,18 +106,12 @@ def _restore_entity(population, directory):
 
 def _restore_holder(simulation, variable, directory):
     storage_dir = os.path.join(directory, variable)
-    is_variable_eternal = simulation.tax_benefit_system.get_variable(variable).definition_period == ETERNITY
-    disk_storage = OnDiskStorage(
-        storage_dir,
-        is_eternal = is_variable_eternal,
-        preserve_storage_dir = True
-        )
-
-    # TODO: decide whether to turn this variable public
-    disk_storage.state = disk_storage.storage.restore(disk_storage.state)
-
+    is_eternal = simulation.tax_benefit_system.get_variable(variable).definition_period == ETERNITY
+    storage = DiskStorage(storage_dir, preserve = True)
+    cache = Cache(storage, is_eternal = is_eternal)
+    cache.state = storage.restore(cache.state)
     holder = simulation.get_holder(variable)
 
-    for period in disk_storage.known_periods():
-        value = disk_storage.get(period)
+    for period in cache.known_periods():
+        value = cache.get(period)
         holder.put_in_cache(value, period)

tests/core/data_storage/test_cache.py

@@ -0,0 +1,117 @@
+import numpy
+
+from openfisca_core import periods
+from openfisca_core.data_storage import Cache, DiskStorage
+
+import pytest
+
+
+@pytest.fixture
+def storage():
+    return DiskStorage(directory = "/tmp/openfisca", preserve = False)
+
+
+@pytest.fixture
+def cache(storage):
+    return Cache(storage = storage)
+
+
+@pytest.fixture
+def period():
+    return periods.period("2020")
+
+
+@pytest.fixture
+def eternal_period():
+    return periods.period(periods.ETERNITY)
+
+
+@pytest.fixture
+def array():
+    return numpy.array([1])
+
+
+def test_get(cache, storage, period, mocker):
+    mocker.patch.object(storage, "get")
+    cache.get(period)
+
+    result = {}, period
+
+    storage.get.assert_called_once_with(*result)
+
+
+def test_get_when_is_eternal(cache, storage, period, eternal_period, mocker):
+    """When it is eternal, input periods are actually ignored."""
+    mocker.patch.object(storage, "get")
+    cache.is_eternal = True
+    cache.get(period)
+
+    result = {}, eternal_period
+
+    storage.get.assert_called_once_with(*result)
+
+
+def test_put(cache, storage, period, array, mocker):
+    mocker.patch.object(storage, "put")
+    cache.put(array, period)
+
+    result = {}, period, array
+
+    storage.put.assert_called_once_with(*result)
+
+
+def test_put_when_is_eternal(cache, storage, period, eternal_period, array, mocker):
+    """When it is eternal, input periods are actually ignored."""
+    mocker.patch.object(storage, "put")
+    cache.is_eternal = True
+    cache.put(array, period)
+
+    result = {}, eternal_period, array
+
+    storage.put.assert_called_once_with(*result)
+
+
+def test_delete(cache, storage, period, mocker):
+    mocker.patch.object(storage, "delete")
+    cache.delete(period)
+
+    result = {}, period
+
+    storage.delete.assert_called_once_with(*result)
+
+
+def test_delete_when_period_is_not_specified(cache, storage, mocker):
+    mocker.patch.object(storage, "delete_all")
+    cache.delete()
+
+    result = {}
+
+    storage.delete_all.assert_called_once_with(result)
+
+
+def test_delete_when_is_eternal(cache, storage, period, eternal_period, mocker):
+    """When it is eternal, input periods are actually ignored."""
+    mocker.patch.object(storage, "delete")
+    cache.is_eternal = True
+    cache.delete(period)
+
+    result = {}, eternal_period
+
+    storage.delete.assert_called_once_with(*result)
+
+
+def test_get_memory_usage(cache, storage, mocker):
+    mocker.patch.object(storage, "memory_usage")
+    cache.memory_usage()
+
+    result = {}
+
+    storage.memory_usage.assert_called_once_with(result)
+
+
+def test_known_periods(cache, period, array):
+    cache.put(array, period)
+
+    result = cache.known_periods()
+
+    assert period in result