docs: add types to holders (1/3)

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+### 43.2.8 [#1285](https://github.com/openfisca/openfisca-core/pull/1285)
+
+#### Documentation
+
+- Add some types to `holders` (1 of 3)
+
 ### 43.2.7 [#1300](https://github.com/openfisca/openfisca-core/pull/1300)
 
 #### Technical changes

openfisca_core/data_storage/__init__.py

@@ -1,7 +1,6 @@
 """Different storage backends for the data of a simulation."""
 
-from . import types
 from .in_memory_storage import InMemoryStorage
 from .on_disk_storage import OnDiskStorage
 
-__all__ = ["InMemoryStorage", "OnDiskStorage", "types"]
+__all__ = ["InMemoryStorage", "OnDiskStorage"]

openfisca_core/data_storage/in_memory_storage.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
 from collections.abc import KeysView, MutableMapping
+from typing import Generic, TypeVar
 
 import numpy
 
-from openfisca_core import periods
-from openfisca_core.periods import DateUnit
+from openfisca_core import periods, types as t
 
-from . import types as t
+#: Type var for numpy arrays (invariant).
+_N = TypeVar("_N", bound=t.VarDType)
 
 
-class InMemoryStorage:
+class InMemoryStorage(Generic[_N]):
     """Storing and retrieving calculated vectors in memory.
 
     Args:
@@ -22,13 +23,13 @@ class InMemoryStorage:
     is_eternal: bool
 
     #: A dictionary containing data that has been stored in memory.
-    _arrays: MutableMapping[t.Period, t.Array[t.DTypeGeneric]]
+    _arrays: MutableMapping[t.Period, t.Array[_N]]
 
     def __init__(self, is_eternal: bool = False) -> None:
         self._arrays = {}
         self.is_eternal = is_eternal
 
-    def get(self, period: None | t.Period = None) -> None | t.Array[t.DTypeGeneric]:
+    def get(self, period: None | t.Period = None) -> None | t.Array[_N]:
         """Retrieve the data for the specified :obj:`.Period` from memory.
 
         Args:
@@ -56,15 +57,15 @@ def get(self, period: None | t.Period = None) -> None | t.Array[t.DTypeGeneric]:
 
         """
         if self.is_eternal:
-            period = periods.period(DateUnit.ETERNITY)
+            period = periods.Period.eternity()
         period = periods.period(period)
 
         values = self._arrays.get(period)
         if values is None:
             return None
         return values
 
-    def put(self, value: t.Array[t.DTypeGeneric], period: None | t.Period) -> None:
+    def put(self, value: t.Array[_N], period: None | t.Period) -> None:
         """Store the specified data in memory for the specified :obj:`.Period`.
 
         Args:
@@ -88,7 +89,7 @@ def put(self, value: t.Array[t.DTypeGeneric], period: None | t.Period) -> None:
 
         """
         if self.is_eternal:
-            period = periods.period(DateUnit.ETERNITY)
+            period = periods.Period.eternity()
         period = periods.period(period)
 
         self._arrays[period] = value
@@ -133,7 +134,7 @@ def delete(self, period: None | t.Period = None) -> None:
             return
 
         if self.is_eternal:
-            period = periods.period(DateUnit.ETERNITY)
+            period = periods.Period.eternity()
         period = periods.period(period)
 
         self._arrays = {

openfisca_core/data_storage/on_disk_storage.py

@@ -1,20 +1,20 @@
 from __future__ import annotations
 
 from collections.abc import KeysView, MutableMapping
+from typing import Generic, TypeVar
 
 import os
 import shutil
 
 import numpy
 
-from openfisca_core import periods
-from openfisca_core.indexed_enums import EnumArray
-from openfisca_core.periods import DateUnit
+from openfisca_core import indexed_enums as enum, periods, types as t
 
-from . import types as t
+#: Type var for numpy arrays (invariant).
+_N = TypeVar("_N", bound=t.VarDType)
 
 
-class OnDiskStorage:
+class OnDiskStorage(Generic[_N]):
     """Storing and retrieving calculated vectors on disk.
 
     Args:
@@ -44,63 +44,24 @@ def __init__(
         storage_dir: str,
         is_eternal: bool = False,
         preserve_storage_dir: bool = False,
-        enums: MutableMapping[str, type[t.Enum]] | None = None,
+        enums: None | MutableMapping[str, type[t.Enum]] = None,
     ) -> None:
         self._files = {}
         self._enums = {} if enums is None else enums
         self.is_eternal = is_eternal
         self.preserve_storage_dir = preserve_storage_dir
         self.storage_dir = storage_dir
 
-    def _decode_file(self, file: str) -> t.Array[t.DTypeGeneric]:
-        """Decode a file by loading its contents as a :mod:`numpy` array.
-
-        Args:
-            file: Path to the file to be decoded.
-
-        Returns:
-            EnumArray: Representing the data in the file.
-            ndarray[generic]: Representing the data in the file.
-
-        Note:
-            If the file is associated with :class:`~indexed_enums.Enum` values, the
-            array is converted back to an :obj:`~indexed_enums.EnumArray` object.
-
-        Examples:
-            >>> import tempfile
-
-            >>> import numpy
-
-            >>> from openfisca_core import data_storage, indexed_enums, periods
-
-            >>> class Housing(indexed_enums.Enum):
-            ...     OWNER = "Owner"
-            ...     TENANT = "Tenant"
-            ...     FREE_LODGER = "Free lodger"
-            ...     HOMELESS = "Homeless"
-
-            >>> array = numpy.array([1])
-            >>> value = indexed_enums.EnumArray(array, Housing)
-            >>> instant = periods.Instant((2017, 1, 1))
-            >>> period = periods.Period(("year", instant, 1))
-
-            >>> with tempfile.TemporaryDirectory() as directory:
-            ...     storage = data_storage.OnDiskStorage(directory)
-            ...     storage.put(value, period)
-            ...     storage._decode_file(storage._files[period])
-            EnumArray([Housing.TENANT])
-
-        """
-        enum = self._enums.get(self.storage_dir)
-
-        if enum is not None:
-            return EnumArray(numpy.load(file), enum)
-
-        array: t.Array[t.DTypeGeneric] = numpy.load(file)
-
-        return array
+    def __del__(self) -> None:
+        if self.preserve_storage_dir:
+            return
+        shutil.rmtree(self.storage_dir)  # Remove the holder temporary files
+        # If the simulation temporary directory is empty, remove it
+        parent_dir = os.path.abspath(os.path.join(self.storage_dir, os.pardir))
+        if not os.listdir(parent_dir):
+            shutil.rmtree(parent_dir)
 
-    def get(self, period: None | t.Period = None) -> None | t.Array[t.DTypeGeneric]:
+    def get(self, period: None | t.Period = None) -> None | t.Array[_N]:
         """Retrieve the data for the specified period from disk.
 
         Args:
@@ -130,15 +91,15 @@ def get(self, period: None | t.Period = None) -> None | t.Array[t.DTypeGeneric]:
 
         """
         if self.is_eternal:
-            period = periods.period(DateUnit.ETERNITY)
+            period = periods.Period.eternity()
         period = periods.period(period)
 
         values = self._files.get(period)
         if values is None:
             return None
         return self._decode_file(values)
 
-    def put(self, value: t.Array[t.DTypeGeneric], period: None | t.Period) -> None:
+    def put(self, value: t.Array[_N], period: None | t.Period) -> None:
         """Store the specified data on disk for the specified period.
 
         Args:
@@ -164,12 +125,12 @@ def put(self, value: t.Array[t.DTypeGeneric], period: None | t.Period) -> None:
 
         """
         if self.is_eternal:
-            period = periods.period(DateUnit.ETERNITY)
+            period = periods.Period.eternity()
         period = periods.period(period)
 
         filename = str(period)
         path = os.path.join(self.storage_dir, filename) + ".npy"
-        if isinstance(value, EnumArray) and value.possible_values is not None:
+        if isinstance(value, enum.EnumArray) and value.possible_values is not None:
             self._enums[self.storage_dir] = value.possible_values
             value = value.view(numpy.ndarray)
         numpy.save(path, value)
@@ -217,7 +178,7 @@ def delete(self, period: None | t.Period = None) -> None:
             return
 
         if self.is_eternal:
-            period = periods.period(DateUnit.ETERNITY)
+            period = periods.Period.eternity()
         period = periods.period(period)
 
         self._files = {
@@ -297,14 +258,53 @@ def restore(self) -> None:
             period = periods.period(filename_core)
             files[period] = path
 
-    def __del__(self) -> None:
-        if self.preserve_storage_dir:
-            return
-        shutil.rmtree(self.storage_dir)  # Remove the holder temporary files
-        # If the simulation temporary directory is empty, remove it
-        parent_dir = os.path.abspath(os.path.join(self.storage_dir, os.pardir))
-        if not os.listdir(parent_dir):
-            shutil.rmtree(parent_dir)
+    def _decode_file(self, file: str) -> t.Array[_N]:
+        """Decode a file by loading its contents as a :mod:`numpy` array.
+
+        Args:
+            file: Path to the file to be decoded.
+
+        Returns:
+            EnumArray: Representing the data in the file.
+            ndarray[generic]: Representing the data in the file.
+
+        Note:
+            If the file is associated with :class:`~indexed_enums.Enum` values, the
+            array is converted back to an :obj:`~indexed_enums.EnumArray` object.
+
+        Examples:
+            >>> import tempfile
+
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, indexed_enums, periods
+
+            >>> class Housing(indexed_enums.Enum):
+            ...     OWNER = "Owner"
+            ...     TENANT = "Tenant"
+            ...     FREE_LODGER = "Free lodger"
+            ...     HOMELESS = "Homeless"
+
+            >>> array = numpy.array([1])
+            >>> value = indexed_enums.EnumArray(array, Housing)
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.put(value, period)
+            ...     storage._decode_file(storage._files[period])
+            EnumArray([Housing.TENANT])
+
+        """
+        enum_class = self._enums.get(self.storage_dir)
+
+        if enum_class is not None:
+            return enum.EnumArray(numpy.load(file), enum_class)
+
+        array: t.Array[_N] = numpy.load(file)
+
+        return array
 
 
 __all__ = ["OnDiskStorage"]

openfisca_core/entities/__init__.py

@@ -1,6 +1,5 @@
 """Provide a way of representing the entities of a rule system."""
 
-from . import types
 from ._core_entity import CoreEntity
 from .entity import Entity
 from .group_entity import GroupEntity
@@ -19,5 +18,4 @@
     "build_entity",
     "check_role_validity",
     "find_role",
-    "types",
 ]

openfisca_core/entities/_core_entity.py

@@ -5,7 +5,8 @@
 import abc
 import os
 
-from . import types as t
+from openfisca_core import types as t
+
 from .role import Role
 
 
@@ -17,8 +18,7 @@ class CoreEntity:
         **__kwargs: Any keyword arguments.
 
     Examples:
-        >>> from openfisca_core import entities
-        >>> from openfisca_core.entities import types as t
+        >>> from openfisca_core import entities, types as t
 
         >>> class Entity(entities.CoreEntity):
         ...     def __init__(self, key):
@@ -61,7 +61,7 @@ def get_variable(
         self,
         variable_name: t.VariableName,
         check_existence: bool = False,
-    ) -> t.Variable | None:
+    ) -> None | t.Variable[t.VarDType]:
         """Get ``variable_name`` from ``variables``.
 
         Args:
@@ -168,7 +168,7 @@ def check_variable_defined_for_entity(self, variable_name: t.VariableName) -> No
 
         """
         entity: None | t.CoreEntity = None
-        variable: None | t.Variable = self.get_variable(
+        variable: None | t.Variable[t.VarDType] = self.get_variable(
             variable_name,
             check_existence=True,
         )

openfisca_core/entities/entity.py

@@ -2,7 +2,8 @@
 
 import textwrap
 
-from . import types as t
+from openfisca_core import types as t
+
 from ._core_entity import CoreEntity