refactor: memory usage (#1238 #1276 #1277)

openfisca_core/holders/__init__.py

@@ -24,11 +24,9 @@
 from . import types
 from .helpers import set_input_dispatch_by_period, set_input_divide_by_period
 from .holder import Holder
-from .memory_usage import MemoryUsage
 
 __all__ = [
     "Holder",
-    "MemoryUsage",
     "set_input_dispatch_by_period",
     "set_input_divide_by_period",
     "types",

openfisca_core/holders/holder.py

@@ -19,7 +19,6 @@
 )
 
 from . import types as t
-from .memory_usage import MemoryUsage
 
 
 class Holder:
@@ -95,7 +94,7 @@ def get_array(self, period):
             return self._disk_storage.get(period)
         return None
 
-    def get_memory_usage(self) -> MemoryUsage:
+    def get_memory_usage(self) -> t.MemoryUsage:
         """Get data about the virtual memory usage of the Holder.
 
         Returns:
@@ -136,7 +135,7 @@ def get_memory_usage(self) -> MemoryUsage:
                'total_nb_bytes': 0...
 
         """
-        usage = MemoryUsage(
+        usage = t.MemoryUsage(
             nb_cells_by_array=self.population.count,
             dtype=self.variable.dtype,
         )

openfisca_core/holders/types.py

@@ -1,3 +1,3 @@
-from openfisca_core.types import CorePopulation, Holder
+from openfisca_core.types import CorePopulation, Holder, MemoryUsage
 
-__all__ = ["CorePopulation", "Holder"]
+__all__ = ["CorePopulation", "Holder", "MemoryUsage"]

openfisca_core/populations/_core_population.py

@@ -27,11 +27,11 @@ class Option(strenum.StrEnum):
 class Calculate(NamedTuple):
     variable: str
     period: t.Period
-    option: Sequence[str] | None
+    option: None | Sequence[str]
 
 
 class MemoryUsageByVariable(TypedDict, total=False):
-    by_variable: dict[str, holders.MemoryUsage]
+    by_variable: dict[str, t.MemoryUsage]
     total_nb_bytes: int
 
 
@@ -186,7 +186,7 @@ def check_array_compatible_with_entity(self, array: t.FloatArray) -> None:
     def check_period_validity(
         self,
         variable_name: str,
-        period: int | str | Period | None,
+        period: None | t.PeriodLike,
     ) -> None:
         if isinstance(period, (int, str, periods.Period)):
             return

openfisca_core/populations/population.py

@@ -12,7 +12,7 @@ class Population(CorePopulation):
     def __init__(self, entity: t.SingleEntity) -> None:
         super().__init__(entity)
 
-    def clone(self, simulation: Simulation) -> t.CorePopulation:
+    def clone(self, simulation: t.Simulation) -> t.CorePopulation:
         result = Population(self.entity)
         result.simulation = simulation
         result._holders = {
@@ -38,7 +38,7 @@ def __getattr__(self, attribute: str) -> projectors.Projector:
     # Helpers
 
     @projectors.projectable
-    def has_role(self, role: Role) -> Array[bool] | None:
+    def has_role(self, role: t.Role) -> None | t.BoolArray:
         """Check if a person has a given role within its `GroupEntity`.
 
         Example:
@@ -63,10 +63,10 @@ def has_role(self, role: Role) -> Array[bool] | None:
     @projectors.projectable
     def value_from_partner(
         self,
-        array: Array[float],
+        array: t.FloatArray,
         entity: projectors.Projector,
-        role: Role,
-    ) -> Array[float] | None:
+        role: t.Role,
+    ) -> None | t.FloatArray:
         self.check_array_compatible_with_entity(array)
         self.entity.check_role_validity(role)
 
@@ -89,9 +89,9 @@ def value_from_partner(
     def get_rank(
         self,
         entity: Population,
-        criteria: Array[float],
+        criteria: t.FloatArray,
         condition: bool = True,
-    ) -> Array[int]:
+    ) -> t.IntArray:
         """Get the rank of a person within an entity according to a criteria.
         The person with rank 0 has the minimum value of criteria.
         If condition is specified, then the persons who don't respect it are not taken into account and their rank is -1.

openfisca_core/populations/types.py

@@ -12,8 +12,10 @@
     EntityKey,
     GroupEntity,
     Holder,
+    MemoryUsage,
     Period,
     PeriodStr,
+    Role,
     Simulation,
     SingleEntity,
     SinglePopulation,
@@ -24,11 +26,15 @@
     bool_ as BoolDType,
     float32 as FloatDType,
     generic as VarDType,
+    int32 as IntDType,
     str_ as StrDType,
 )
 
 # Commons
 
+#: Type alias for an array of strings.
+IntArray: TypeAlias = Array[IntDType]
+
 #: Type alias for an array of strings.
 StrArray: TypeAlias = Array[StrDType]
 
@@ -63,7 +69,9 @@
     "EntityKey",
     "GroupEntity",
     "Holder",
+    "MemoryUsage",
     "Period",
+    "Role",
     "Simulation",
     "SingleEntity",
     "SinglePopulation",

openfisca_core/types.py

@@ -3,7 +3,7 @@
 from collections.abc import Iterable, Mapping, Sequence, Sized
 from numpy.typing import DTypeLike, NDArray
 from typing import NewType, TypeVar, Union
-from typing_extensions import Protocol, Self, TypeAlias
+from typing_extensions import Protocol, Self, TypeAlias, TypedDict
 
 import abc
 import enum
@@ -136,6 +136,16 @@ def clone(self, population: CorePopulation, /) -> Holder: ...
     def get_memory_usage(self, /) -> Mapping[str, object]: ...
 
 
+class MemoryUsage(TypedDict, total=False):
+    cell_size: int
+    dtype: DTypeLike
+    nb_arrays: int
+    nb_cells_by_array: int
+    nb_requests: int
+    nb_requests_by_array: int
+    total_nb_bytes: int
+
+
 # Parameters