Merge branch 'master' into ci/fix-conda-publish

CHANGELOG.md

@@ -1,5 +1,23 @@
 # Changelog
 
+### 42.0.7 [#1264](https://github.com/openfisca/openfisca-core/pull/1264)
+
+#### Technical changes
+
+- Add typing to `data_storage` module
+
+### 42.0.6 [#1263](https://github.com/openfisca/openfisca-core/pull/1263)
+
+#### Documentation
+
+- Fix docs of the `data_storage` module
+
+### 42.0.5 [#1261](https://github.com/openfisca/openfisca-core/pull/1261)
+
+#### Technical changes
+
+- Fix doctests of `data_storage` module
+
 ### 42.0.4 [#1257](https://github.com/openfisca/openfisca-core/pull/1257)
 
 #### Technical changes

openfisca_core/data_storage/__init__.py

@@ -21,5 +21,8 @@
 #
 # See: https://www.python.org/dev/peps/pep-0008/#imports
 
-from .in_memory_storage import InMemoryStorage  # noqa: F401
-from .on_disk_storage import OnDiskStorage  # noqa: F401
+from . import types
+from .in_memory_storage import InMemoryStorage
+from .on_disk_storage import OnDiskStorage
+
+__all__ = ["InMemoryStorage", "OnDiskStorage", "types"]

openfisca_core/data_storage/in_memory_storage.py

@@ -1,17 +1,59 @@
+from __future__ import annotations
+
+from collections.abc import KeysView, MutableMapping
+
 import numpy
 
 from openfisca_core import periods
 from openfisca_core.periods import DateUnit
 
+from . import types as t
+
 
 class InMemoryStorage:
-    """Low-level class responsible for storing and retrieving calculated vectors in memory."""
+    """Storing and retrieving calculated vectors in memory.
+
+    Args:
+        is_eternal: Whether the storage is eternal.
+
+    """
+
+    #: Whether the storage is eternal.
+    is_eternal: bool
+
+    #: A dictionary containing data that has been stored in memory.
+    _arrays: MutableMapping[t.Period, t.Array[t.DTypeGeneric]]
 
-    def __init__(self, is_eternal=False) -> None:
+    def __init__(self, is_eternal: bool = False) -> None:
         self._arrays = {}
         self.is_eternal = is_eternal
 
-    def get(self, period):
+    def get(self, period: None | t.Period = None) -> None | t.Array[t.DTypeGeneric]:
+        """Retrieve the data for the specified period from memory.
+
+        Args:
+            period: The period for which data should be retrieved.
+
+        Returns:
+            The data for the specified period, or None if no data is available.
+
+        Examples:
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, periods
+
+            >>> storage = data_storage.InMemoryStorage()
+            >>> value = numpy.array([1, 2, 3])
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> storage.put(value, period)
+
+            >>> storage.get(period)
+            array([1, 2, 3])
+
+        """
+
         if self.is_eternal:
             period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
@@ -21,14 +63,72 @@ def get(self, period):
             return None
         return values
 
-    def put(self, value, period) -> None:
+    def put(self, value: t.Array[t.DTypeGeneric], period: None | t.Period) -> None:
+        """Store the specified data in memory for the specified period.
+
+        Args:
+            value: The data to store
+            period: The period for which the data should be stored.
+
+        Examples:
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, periods
+
+            >>> storage = data_storage.InMemoryStorage()
+            >>> value = numpy.array([1, "2", "salary"])
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> storage.put(value, period)
+
+            >>> storage.get(period)
+            array(['1', '2', 'salary'], dtype='<U21')
+
+        """
+
         if self.is_eternal:
             period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
 
         self._arrays[period] = value
 
-    def delete(self, period=None) -> None:
+    def delete(self, period: None | t.Period = None) -> None:
+        """Delete the data for the specified period from memory.
+
+        Args:
+            period: The period for which data should be deleted.
+
+        Note:
+            If ``period`` is specified, all data will be deleted.
+
+        Examples:
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, periods
+
+            >>> storage = data_storage.InMemoryStorage()
+            >>> value = numpy.array([1, 2, 3])
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> storage.put(value, period)
+
+            >>> storage.get(period)
+            array([1, 2, 3])
+
+            >>> storage.delete(period)
+
+            >>> storage.get(period)
+
+            >>> storage.put(value, period)
+
+            >>> storage.delete()
+
+            >>> storage.get(period)
+
+        """
+
         if period is None:
             self._arrays = {}
             return
@@ -43,10 +143,45 @@ def delete(self, period=None) -> None:
             if not period.contains(period_item)
         }
 
-    def get_known_periods(self):
+    def get_known_periods(self) -> KeysView[t.Period]:
+        """List of storage's known periods.
+
+        Returns:
+            A sequence containing the storage's known periods.
+
+        Examples:
+            >>> from openfisca_core import data_storage, periods
+
+            >>> storage = data_storage.InMemoryStorage()
+            >>> storage.get_known_periods()
+            dict_keys([])
+
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+            >>> storage.put([], period)
+
+            >>> storage.get_known_periods()
+            dict_keys([Period(('year', Instant((2017, 1, 1)), 1))])
+
+        """
+
         return self._arrays.keys()
 
-    def get_memory_usage(self):
+    def get_memory_usage(self) -> t.MemoryUsage:
+        """Memory usage of the storage.
+
+        Returns:
+            A dictionary representing the storage's memory usage.
+
+        Examples:
+            >>> from openfisca_core import data_storage
+
+            >>> storage = data_storage.InMemoryStorage()
+            >>> storage.get_memory_usage()
+            {'nb_arrays': 0, 'total_nb_bytes': 0, 'cell_size': nan}
+
+        """
+
         if not self._arrays:
             return {
                 "nb_arrays": 0,
@@ -61,3 +196,6 @@ def get_memory_usage(self):
             "total_nb_bytes": array.nbytes * nb_arrays,
             "cell_size": array.itemsize,
         }
+
+
+__all__ = ["InMemoryStorage"]