Bump version

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+### 35.11.3 [#1167](https://github.com/openfisca/openfisca-core/pull/1167)
+
+#### Technical changes
+
+- Add typing to `data_storage`.
+
 ### 35.11.2 [#1166](https://github.com/openfisca/openfisca-core/pull/1166)
 
 #### Technical changes

openfisca_core/data_storage/_arrays.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from typing import Dict
+
+import collections
+import dataclasses
+
+import numpy
+
+from openfisca_core import types
+
+
+@dataclasses.dataclass(frozen = True)
+class Arrays(collections.UserDict):
+    data: Dict[types.Period, numpy.ndarray]

openfisca_core/data_storage/_enums.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Dict, Type
+
+import collections
+import dataclasses
+
+from openfisca_core import types
+
+
+@dataclasses.dataclass(frozen = True)
+class Enums(collections.UserDict):
+    data: Dict[str, Type[types.Enum]]

openfisca_core/data_storage/_files.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Dict
+
+import collections
+import dataclasses
+
+from openfisca_core import types
+
+
+@dataclasses.dataclass(frozen = True)
+class Files(collections.UserDict):
+    data: Dict[types.Period, str]

openfisca_core/data_storage/in_memory_storage.py

@@ -1,21 +1,23 @@
 from __future__ import annotations
 
-from typing import Any, Dict, KeysView, Optional
+from typing import Any, Optional, Sequence
 
 import numpy
 
 from openfisca_core import periods, types
 
+from ._arrays import Arrays
+
 
 class InMemoryStorage:
     """
     Low-level class responsible for storing and retrieving calculated vectors in memory
     """
 
-    _arrays: Dict[types.Period, numpy.ndarray]
+    _arrays: Arrays
 
     def __init__(self, is_eternal: bool = False) -> None:
-        self._arrays = {}
+        self._arrays = Arrays({})
         self.is_eternal = is_eternal
 
     def get(self, period: types.Period) -> Any:
@@ -32,40 +34,46 @@ def get(self, period: types.Period) -> Any:
     def put(self, value: Any, period: types.Period) -> None:
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
-        period = periods.period(period)
 
-        self._arrays[period] = value
+        else:
+            period = periods.period(period)
+
+        self._arrays = Arrays({period: value, **self._arrays})
 
     def delete(self, period: Optional[types.Period] = None) -> None:
         if period is None:
-            self._arrays = {}
-            return
+            self._arrays = Arrays({})
+            return None
 
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
-        period = periods.period(period)
 
-        self._arrays = {
+        else:
+            period = periods.period(period)
+
+        self._arrays = Arrays({
             period_item: value
             for period_item, value in self._arrays.items()
             if not period.contains(period_item)
-            }
+            })
 
-    def get_known_periods(self) -> KeysView[types.Period]:
-        return self._arrays.keys()
+    def get_known_periods(self) -> Sequence[types.Period]:
+        return list(self._arrays.keys())
 
     def get_memory_usage(self) -> types.MemoryUsage:
         if not self._arrays:
-            return dict(
+            return types.MemoryUsage(
+                cell_size = numpy.nan,
                 nb_arrays = 0,
                 total_nb_bytes = 0,
-                cell_size = numpy.nan,
                 )
 
         nb_arrays = len(self._arrays)
         array = next(iter(self._arrays.values()))
-        return dict(
-            nb_arrays = nb_arrays,
-            total_nb_bytes = array.nbytes * nb_arrays,
+        total = array.nbytes * nb_arrays
+
+        return types.MemoryUsage(
             cell_size = array.itemsize,
+            nb_arrays = nb_arrays,
+            total_nb_bytes = total,
             )

openfisca_core/data_storage/on_disk_storage.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, KeysView, NoReturn, Optional, Type
+from typing import Any, NoReturn, Optional, Sequence
 
 import os
 import shutil
@@ -9,35 +9,53 @@
 
 from openfisca_core import periods, indexed_enums as enums, types
 
+from ._enums import Enums
+from ._files import Files
+
 
 class OnDiskStorage:
-    """
-    Low-level class responsible for storing and retrieving calculated vectors on disk
+    """Class responsible for storing and retrieving calculated vectors on disk.
+
+    Attributes:
+        _enums: ?
+        _files: ?
+        is_eternal: ?
+        storage_dir: Path to store calculated vectors.
+        preserve_storage_dir: ?
+
+    Args:
+        storage_dir: Path to store calculated vectors.
+        is_eternal: ?
+        preserve_storage_dir: ?
+
     """
 
-    _files: Dict[types.Period, str]
-    _enums: Dict[str, Type[enums.Enum]]
+    _enums: Enums
+    _files: Files
+    is_eternal: bool
+    storage_dir: str
+    preserve_storage_dir: bool
 
     def __init__(
             self,
             storage_dir: str,
             is_eternal: bool = False,
             preserve_storage_dir: bool = False,
             ) -> None:
-        self._files = {}
-        self._enums = {}
+        self._enums = Enums({})
+        self._files = Files({})
         self.is_eternal = is_eternal
-        self.preserve_storage_dir = preserve_storage_dir
         self.storage_dir = storage_dir
+        self.preserve_storage_dir = preserve_storage_dir
 
     def _decode_file(self, file: str) -> Any:
-        enum: Optional[Type[enums.Enum]]
         enum = self._enums.get(file)
+        load = numpy.load(file)
 
-        if enum is not None:
-            return enums.EnumArray(numpy.load(file), enum)
-        else:
-            return numpy.load(file)
+        if enum is None:
+            return load
+
+        return enums.EnumArray(load, enum)
 
     def get(self, period: types.Period) -> Any:
         if self.is_eternal:
@@ -53,48 +71,52 @@ def get(self, period: types.Period) -> Any:
     def put(self, value: Any, period: types.Period) -> None:
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
-        period = periods.period(period)
+
+        else:
+            period = periods.period(period)
 
         filename = str(period)
-        path = os.path.join(self.storage_dir, filename) + '.npy'
+        path = os.path.join(self.storage_dir, filename) + ".npy"
+
         if isinstance(value, enums.EnumArray):
-            self._enums[path] = value.possible_values
+            self._enums = Enums({path: value.possible_values, **self._enums})
             value = value.view(numpy.ndarray)
+
         numpy.save(path, value)
-        self._files[period] = path
+        self._files = Files({period: path, **self._files})
 
     def delete(self, period: Optional[types.Period] = None) -> None:
         if period is None:
-            self._files = {}
+            self._files = Files({})
             return
 
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
         period = periods.period(period)
 
         if period is not None:
-            self._files = {
+            self._files = Files({
                 period_item: value
                 for period_item, value in self._files.items()
                 if not period.contains(period_item)
-                }
+                })
 
-    def get_known_periods(self) -> KeysView[types.Period]:
-        return self._files.keys()
+    def get_known_periods(self) -> Sequence[types.Period]:
+        return list(self._files.keys())
 
     def get_memory_usage(self) -> NoReturn:
         raise NotImplementedError
 
     def restore(self) -> None:
-        self._files = {}
+        self._files = Files({})
         # Restore self._files from content of storage_dir.
         for filename in os.listdir(self.storage_dir):
             if not filename.endswith('.npy'):
                 continue
             path = os.path.join(self.storage_dir, filename)
             filename_core = filename.rsplit('.', 1)[0]
             period = periods.period(filename_core)
-            self._files[period] = path
+            self._files = Files({period: path, **self._files})
 
     def __del__(self) -> None:
         if self.preserve_storage_dir:

openfisca_core/indexed_enums/enum_array.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
 
-import typing
 from typing import Any, NoReturn, Optional, Type
 
 import numpy
 
-if typing.TYPE_CHECKING:
-    from openfisca_core.indexed_enums import Enum
+from openfisca_core import types
 
 
 class EnumArray(numpy.ndarray):
@@ -22,7 +20,7 @@ class EnumArray(numpy.ndarray):
     def __new__(
             cls,
             input_array: numpy.int_,
-            possible_values: Optional[Type[Enum]] = None,
+            possible_values: Optional[Type[types.Enum]] = None,
             ) -> EnumArray:
         obj = numpy.asarray(input_array).view(cls)
         obj.possible_values = possible_values

openfisca_core/types/__init__.py

@@ -8,6 +8,8 @@
     * :attr:`.Array`
     * ``ArrayLike``
     * :attr:`.Cache`
+    * :attr:`.Enum`
+    * :attr:`.EnumArray`
     * :attr:`.Entity`
     * :attr:`.Formula`
     * :attr:`.Holder`
@@ -54,6 +56,8 @@
 from ._data import (  # noqa: F401
     Array,
     ArrayLike,
+    Enum,
+    EnumArray,
     Instant,
     MemoryUsage,
     Period,
@@ -80,6 +84,8 @@
     "Array",
     "ArrayLike",
     "Entity",
+    "Enum",
+    "EnumArray",
     "Formula",
     "Holder",
     "Instant",

openfisca_core/types/_data.py

@@ -12,6 +12,11 @@
     >>> this == that
     True
 
+    >>> that = periods.Instant((1234, 7, 8))
+    >>> this == that
+    False
+
+
 """
 
 from __future__ import annotations
@@ -72,6 +77,14 @@
 """
 
 
+class Enum(Protocol):
+    """Enum protocol."""
+
+
+class EnumArray(Protocol):
+    """EnumArray protocol."""
+
+
 class Instant(Protocol):
     """Instant protocol."""
 

openfisca_core/types/_domain.py

@@ -6,13 +6,31 @@
 if the data they hold is different, they are equal but not fungible.
 
 Examples:
+    If we take entities, they are equal as long as they share the same ``key``.
+    Let's take the following example:
+
     >>> from openfisca_core import entities
 
     >>> this = entities.Entity(1, "a", "b", "c")
     >>> that = entities.Entity(1, "d", "f", "g")
     >>> this == that
     True
 
+    As you can see, ``this`` and ``that`` are equal because they share the same
+    ``key``:
+
+    >>> this.key == that.key
+    True
+
+    The opposite is also true:
+
+    >>> that = entities.Entity(2, "a", "b", "c")
+    >>> this == that
+    False
+
+    >>> this.key == that.key
+    False
+
 """
 
 from __future__ import annotations

openfisca_tasks/lint.mk

@@ -42,7 +42,6 @@ check-types:
 ## Run static type checkers for type errors (strict).
 lint-typing-strict: \
 	lint-typing-strict-commons \
-	lint-typing-strict-data_storage \
 	lint-typing-strict-types \
 	;