OpenFisca-Core-test 46.0.0__py3-none-any.whl

openfisca_core/holders/holder.py

@@ -0,0 +1,328 @@
+from __future__ import annotations
+
+import os
+import warnings
+from collections.abc import Sequence
+from typing import Any
+
+import numpy
+import psutil
+
+from openfisca_core import (
+    commons,
+    errors,
+    periods,
+    types,
+)
+from openfisca_core import (
+    data_storage as storage,
+)
+from openfisca_core import (
+    indexed_enums as enums,
+)
+
+from . import types as t
+
+
+class Holder:
+    """A holder keeps tracks of a variable values after they have been calculated, or set as an input."""
+
+    def __init__(self, variable, population) -> None:
+        self.population = population
+        self.variable = variable
+        self.simulation = population.simulation
+        self._eternal = self.variable.definition_period == periods.DateUnit.ETERNITY
+        self._memory_storage = storage.InMemoryStorage(is_eternal=self._eternal)
+
+        # By default, do not activate on-disk storage, or variable dropping
+        self._disk_storage = None
+        self._on_disk_storable = False
+        self._do_not_store = False
+        if self.simulation and self.simulation.memory_config:
+            if (
+                self.variable.name
+                not in self.simulation.memory_config.priority_variables
+            ):
+                self._disk_storage = self.create_disk_storage()
+                self._on_disk_storable = True
+            if self.variable.name in self.simulation.memory_config.variables_to_drop:
+                self._do_not_store = True
+
+    def clone(self, population: t.CorePopulation) -> t.Holder:
+        """Copy the holder just enough to be able to run a new simulation without modifying the original simulation."""
+        new = commons.empty_clone(self)
+        new_dict = new.__dict__
+
+        for key, value in self.__dict__.items():
+            if key not in ("population", "formula", "simulation"):
+                new_dict[key] = value
+
+        new_dict["population"] = population
+        new_dict["simulation"] = population.simulation
+
+        return new
+
+    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 storage.OnDiskStorage(
+            storage_dir,
+            self._eternal,
+            preserve_storage_dir=preserve,
+        )
+
+    def delete_arrays(self, period=None) -> None:
+        """If ``period`` is ``None``, remove all known values of the variable.
+
+        If ``period`` is not ``None``, only remove all values for any period included in period (e.g. if period is "2017", values for "2017-01", "2017-07", etc. would be removed)
+        """
+        self._memory_storage.delete(period)
+        if self._disk_storage:
+            self._disk_storage.delete(period)
+
+    def get_array(self, period):
+        """Get the value of the variable for the given period.
+
+        If the value is not known, return ``None``.
+        """
+        if self.variable.is_neutralized:
+            return self.default_array()
+        value = self._memory_storage.get(period)
+        if value is not None:
+            return value
+        if self._disk_storage:
+            return self._disk_storage.get(period)
+        return None
+
+    def get_memory_usage(self) -> t.MemoryUsage:
+        """Get data about the virtual memory usage of the Holder.
+
+        Returns:
+            Memory usage data.
+
+        Examples:
+            >>> from pprint import pprint
+
+            >>> from openfisca_core import (
+            ...     entities,
+            ...     populations,
+            ...     simulations,
+            ...     taxbenefitsystems,
+            ...     variables,
+            ... )
+
+            >>> entity = entities.Entity("", "", "", "")
+
+            >>> class MyVariable(variables.Variable):
+            ...     definition_period = periods.DateUnit.YEAR
+            ...     entity = entity
+            ...     value_type = int
+
+            >>> population = populations.Population(entity)
+            >>> variable = MyVariable()
+            >>> holder = Holder(variable, population)
+
+            >>> tbs = taxbenefitsystems.TaxBenefitSystem([entity])
+            >>> entities = {entity.key: population}
+            >>> simulation = simulations.Simulation(tbs, entities)
+            >>> holder.simulation = simulation
+
+            >>> pprint(holder.get_memory_usage(), indent=3)
+            {  'cell_size': nan,
+               'dtype': <class 'numpy.int32'>,
+               'nb_arrays': 0,
+               'nb_cells_by_array': 0,
+               'total_nb_bytes': 0...
+
+        """
+        usage = t.MemoryUsage(
+            nb_cells_by_array=self.population.count,
+            dtype=self.variable.dtype,
+        )
+
+        usage.update(self._memory_storage.get_memory_usage())
+
+        if self.simulation.trace:
+            nb_requests = self.simulation.tracer.get_nb_requests(self.variable.name)
+            usage.update(
+                {
+                    "nb_requests": nb_requests,
+                    "nb_requests_by_array": (
+                        nb_requests / float(usage["nb_arrays"])
+                        if usage["nb_arrays"] > 0
+                        else numpy.nan
+                    ),
+                },
+            )
+
+        return usage
+
+    def get_known_periods(self):
+        """Get the list of periods the variable value is known for."""
+        return list(self._memory_storage.get_known_periods()) + list(
+            self._disk_storage.get_known_periods() if self._disk_storage else [],
+        )
+
+    def set_input(
+        self,
+        period: types.Period,
+        array: numpy.ndarray | Sequence[Any],
+    ) -> numpy.ndarray | None:
+        """Set a Variable's array of values of a given Period.
+
+        Args:
+            period: The period at which the value is set.
+            array: The input value for the variable.
+
+        Returns:
+            The set input array.
+
+        Note:
+            If a ``set_input`` property has been set for the variable, this
+            method may accept inputs for periods not matching the
+            ``definition_period`` of the Variable. To read
+            more about this, check the `documentation`_.
+
+        Examples:
+            >>> from openfisca_core import entities, populations, variables
+
+            >>> entity = entities.Entity("", "", "", "")
+
+            >>> class MyVariable(variables.Variable):
+            ...     definition_period = periods.DateUnit.YEAR
+            ...     entity = entity
+            ...     value_type = float
+
+            >>> variable = MyVariable()
+
+            >>> population = populations.Population(entity)
+            >>> population.count = 2
+
+            >>> holder = Holder(variable, population)
+            >>> holder.set_input("2018", numpy.array([12.5, 14]))
+            >>> holder.get_array("2018")
+            array([12.5, 14. ], dtype=float32)
+
+            >>> holder.set_input("2018", [12.5, 14])
+            >>> holder.get_array("2018")
+            array([12.5, 14. ], dtype=float32)
+
+        .. _documentation:
+            https://openfisca.org/doc/coding-the-legislation/35_periods.html#set-input-automatically-process-variable-inputs-defined-for-periods-not-matching-the-definition-period
+
+        """
+        period = periods.period(period)
+
+        if period.unit == periods.DateUnit.ETERNITY and not self._eternal:
+            error_message = os.linesep.join(
+                [
+                    "Unable to set a value for variable {1} for {0}.",
+                    "{1} is only defined for {2}s. Please adapt your input.",
+                ],
+            ).format(
+                periods.DateUnit.ETERNITY.upper(),
+                self.variable.name,
+                self.variable.definition_period,
+            )
+            raise errors.PeriodMismatchError(
+                self.variable.name,
+                period,
+                self.variable.definition_period,
+                error_message,
+            )
+        if self.variable.is_neutralized:
+            warning_message = f"You cannot set a value for the variable {self.variable.name}, as it has been neutralized. The value you provided ({array}) will be ignored."
+            return warnings.warn(warning_message, Warning, stacklevel=2)
+        if self.variable.value_type in (float, int) and isinstance(array, str):
+            array = commons.eval_expression(array)
+        if self.variable.set_input:
+            return self.variable.set_input(self, period, array)
+        return self._set(period, array)
+
+    def _to_array(self, value):
+        if not isinstance(value, numpy.ndarray):
+            value = numpy.asarray(value)
+        if value.ndim == 0:
+            # 0-dim arrays are casted to scalar when they interact with float. We don't want that.
+            value = value.reshape(1)
+        if len(value) != self.population.count:
+            msg = f'Unable to set value "{value}" for variable "{self.variable.name}", as its length is {len(value)} while there are {self.population.count} {self.population.entity.plural} in the simulation.'
+            raise ValueError(
+                msg,
+            )
+        if self.variable.value_type == enums.Enum:
+            value = self.variable.possible_values.encode(value)
+        if value.dtype != self.variable.dtype:
+            try:
+                value = value.astype(self.variable.dtype)
+            except ValueError:
+                msg = f'Unable to set value "{value}" for variable "{self.variable.name}", as the variable dtype "{self.variable.dtype}" does not match the value dtype "{value.dtype}".'
+                raise ValueError(
+                    msg,
+                )
+        return value
+
+    def _set(self, period, value) -> None:
+        value = self._to_array(value)
+        if not self._eternal:
+            if period is None:
+                msg = (
+                    f"A period must be specified to set values, except for variables with "
+                    f"{periods.DateUnit.ETERNITY.upper()} as as period_definition."
+                )
+                raise ValueError(
+                    msg,
+                )
+            if self.variable.definition_period != period.unit or period.size > 1:
+                name = self.variable.name
+                period_size_adj = (
+                    f"{period.unit}"
+                    if (period.size == 1)
+                    else f"{period.size}-{period.unit}s"
+                )
+                error_message = os.linesep.join(
+                    [
+                        f'Unable to set a value for variable "{name}" for {period_size_adj}-long period "{period}".',
+                        f'"{name}" can only be set for one {self.variable.definition_period} at a time. Please adapt your input.',
+                        f'If you are the maintainer of "{name}", you can consider adding it a set_input attribute to enable automatic period casting.',
+                    ],
+                )
+
+                raise errors.PeriodMismatchError(
+                    self.variable.name,
+                    period,
+                    self.variable.definition_period,
+                    error_message,
+                )
+
+        should_store_on_disk = (
+            self._on_disk_storable
+            and self._memory_storage.get(period) is None
+            and psutil.virtual_memory().percent  # If there is already a value in memory, replace it and don't put a new value in the disk storage
+            >= self.simulation.memory_config.max_memory_occupation_pc
+        )
+
+        if should_store_on_disk:
+            self._disk_storage.put(value, period)
+        else:
+            self._memory_storage.put(value, period)
+
+    def put_in_cache(self, value, period) -> None:
+        if self._do_not_store:
+            return
+
+        if (
+            self.simulation.opt_out_cache
+            and self.simulation.tax_benefit_system.cache_blacklist
+            and self.variable.name in self.simulation.tax_benefit_system.cache_blacklist
+        ):
+            return
+
+        self._set(period, value)
+
+    def default_array(self):
+        """Return a new array of the appropriate length for the entity, filled with the variable default values."""
+        return self.variable.default_array(self.population.count)

openfisca_core/holders/tests/test_helpers.py

@@ -0,0 +1,134 @@
+import pytest
+
+from openfisca_core import holders, tools
+from openfisca_core.entities import Entity
+from openfisca_core.holders import Holder
+from openfisca_core.periods import DateUnit, Instant, Period
+from openfisca_core.populations import Population
+from openfisca_core.variables import Variable
+
+
+@pytest.fixture
+def people():
+    return Entity(
+        key="person",
+        plural="people",
+        label="An individual member of a larger group.",
+        doc="People have the particularity of not being someone else.",
+    )
+
+
+@pytest.fixture
+def Income(people):
+    return type(
+        "Income",
+        (Variable,),
+        {"value_type": float, "entity": people},
+    )
+
+
+@pytest.fixture
+def population(people):
+    population = Population(people)
+    population.count = 1
+    return population
+
+
+@pytest.mark.parametrize(
+    ("dispatch_unit", "definition_unit", "values", "expected"),
+    [
+        (DateUnit.YEAR, DateUnit.YEAR, [1.0], [3.0]),
+        (DateUnit.YEAR, DateUnit.MONTH, [1.0], [36.0]),
+        (DateUnit.YEAR, DateUnit.DAY, [1.0], [1096.0]),
+        (DateUnit.YEAR, DateUnit.WEEK, [1.0], [157.0]),
+        (DateUnit.YEAR, DateUnit.WEEKDAY, [1.0], [1096.0]),
+        (DateUnit.MONTH, DateUnit.YEAR, [1.0], [1.0]),
+        (DateUnit.MONTH, DateUnit.MONTH, [1.0], [3.0]),
+        (DateUnit.MONTH, DateUnit.DAY, [1.0], [90.0]),
+        (DateUnit.MONTH, DateUnit.WEEK, [1.0], [13.0]),
+        (DateUnit.MONTH, DateUnit.WEEKDAY, [1.0], [90.0]),
+        (DateUnit.DAY, DateUnit.YEAR, [1.0], [1.0]),
+        (DateUnit.DAY, DateUnit.MONTH, [1.0], [1.0]),
+        (DateUnit.DAY, DateUnit.DAY, [1.0], [3.0]),
+        (DateUnit.DAY, DateUnit.WEEK, [1.0], [1.0]),
+        (DateUnit.DAY, DateUnit.WEEKDAY, [1.0], [3.0]),
+        (DateUnit.WEEK, DateUnit.YEAR, [1.0], [1.0]),
+        (DateUnit.WEEK, DateUnit.MONTH, [1.0], [1.0]),
+        (DateUnit.WEEK, DateUnit.DAY, [1.0], [21.0]),
+        (DateUnit.WEEK, DateUnit.WEEK, [1.0], [3.0]),
+        (DateUnit.WEEK, DateUnit.WEEKDAY, [1.0], [21.0]),
+        (DateUnit.WEEKDAY, DateUnit.YEAR, [1.0], [1.0]),
+        (DateUnit.WEEKDAY, DateUnit.MONTH, [1.0], [1.0]),
+        (DateUnit.WEEKDAY, DateUnit.DAY, [1.0], [3.0]),
+        (DateUnit.WEEKDAY, DateUnit.WEEK, [1.0], [1.0]),
+        (DateUnit.WEEKDAY, DateUnit.WEEKDAY, [1.0], [3.0]),
+    ],
+)
+def test_set_input_dispatch_by_period(
+    Income,
+    population,
+    dispatch_unit,
+    definition_unit,
+    values,
+    expected,
+) -> None:
+    Income.definition_period = definition_unit
+    income = Income()
+    holder = Holder(income, population)
+    instant = Instant((2022, 1, 1))
+    dispatch_period = Period((dispatch_unit, instant, 3))
+
+    holders.set_input_dispatch_by_period(holder, dispatch_period, values)
+    total = sum(map(holder.get_array, holder.get_known_periods()))
+
+    tools.assert_near(total, expected, absolute_error_margin=0.001)
+
+
+@pytest.mark.parametrize(
+    ("divide_unit", "definition_unit", "values", "expected"),
+    [
+        (DateUnit.YEAR, DateUnit.YEAR, [3.0], [1.0]),
+        (DateUnit.YEAR, DateUnit.MONTH, [36.0], [1.0]),
+        (DateUnit.YEAR, DateUnit.DAY, [1095.0], [1.0]),
+        (DateUnit.YEAR, DateUnit.WEEK, [157.0], [1.0]),
+        (DateUnit.YEAR, DateUnit.WEEKDAY, [1095.0], [1.0]),
+        (DateUnit.MONTH, DateUnit.YEAR, [1.0], [1.0]),
+        (DateUnit.MONTH, DateUnit.MONTH, [3.0], [1.0]),
+        (DateUnit.MONTH, DateUnit.DAY, [90.0], [1.0]),
+        (DateUnit.MONTH, DateUnit.WEEK, [13.0], [1.0]),
+        (DateUnit.MONTH, DateUnit.WEEKDAY, [90.0], [1.0]),
+        (DateUnit.DAY, DateUnit.YEAR, [1.0], [1.0]),
+        (DateUnit.DAY, DateUnit.MONTH, [1.0], [1.0]),
+        (DateUnit.DAY, DateUnit.DAY, [3.0], [1.0]),
+        (DateUnit.DAY, DateUnit.WEEK, [1.0], [1.0]),
+        (DateUnit.DAY, DateUnit.WEEKDAY, [3.0], [1.0]),
+        (DateUnit.WEEK, DateUnit.YEAR, [1.0], [1.0]),
+        (DateUnit.WEEK, DateUnit.MONTH, [1.0], [1.0]),
+        (DateUnit.WEEK, DateUnit.DAY, [21.0], [1.0]),
+        (DateUnit.WEEK, DateUnit.WEEK, [3.0], [1.0]),
+        (DateUnit.WEEK, DateUnit.WEEKDAY, [21.0], [1.0]),
+        (DateUnit.WEEKDAY, DateUnit.YEAR, [1.0], [1.0]),
+        (DateUnit.WEEKDAY, DateUnit.MONTH, [1.0], [1.0]),
+        (DateUnit.WEEKDAY, DateUnit.DAY, [3.0], [1.0]),
+        (DateUnit.WEEKDAY, DateUnit.WEEK, [1.0], [1.0]),
+        (DateUnit.WEEKDAY, DateUnit.WEEKDAY, [3.0], [1.0]),
+    ],
+)
+def test_set_input_divide_by_period(
+    Income,
+    population,
+    divide_unit,
+    definition_unit,
+    values,
+    expected,
+) -> None:
+    Income.definition_period = definition_unit
+    income = Income()
+    holder = Holder(income, population)
+    instant = Instant((2022, 1, 1))
+    divide_period = Period((divide_unit, instant, 3))
+
+    holders.set_input_divide_by_period(holder, divide_period, values)
+    last = holder.get_array(holder.get_known_periods()[-1])
+
+    tools.assert_near(last, expected, absolute_error_margin=0.001)

openfisca_core/holders/types.py

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

openfisca_core/indexed_enums/__init__.py

@@ -0,0 +1,18 @@
+"""Enumerations for variables with a limited set of possible values."""
+
+from . import types
+from ._enum_type import EnumType
+from ._errors import EnumEncodingError, EnumMemberNotFoundError
+from .config import ENUM_ARRAY_DTYPE
+from .enum import Enum
+from .enum_array import EnumArray
+
+__all__ = [
+    "ENUM_ARRAY_DTYPE",
+    "Enum",
+    "EnumArray",
+    "EnumEncodingError",
+    "EnumMemberNotFoundError",
+    "EnumType",
+    "types",
+]

openfisca_core/indexed_enums/_enum_type.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import final
+
+import numpy
+
+from . import types as t
+
+
+@final
+class EnumType(t.EnumType):
+    """Meta class for creating an indexed :class:`.Enum`.
+
+    Examples:
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Enum(enum.Enum, metaclass=enum.EnumType):
+        ...     pass
+
+        >>> Enum.items
+        Traceback (most recent call last):
+        AttributeError: ...
+
+        >>> class Housing(Enum):
+        ...     OWNER = "Owner"
+        ...     TENANT = "Tenant"
+
+        >>> Housing.indices
+        array([0, 1], dtype=int16)
+
+        >>> Housing.names
+        array(['OWNER', 'TENANT'], dtype='<U6')
+
+        >>> Housing.enums
+        array([Housing.OWNER, Housing.TENANT], dtype=object)
+
+    """
+
+    def __new__(
+        metacls,
+        name: str,
+        bases: tuple[type, ...],
+        classdict: t.EnumDict,
+        **kwds: object,
+    ) -> t.EnumType:
+        """Create a new indexed enum class."""
+        # Create the enum class.
+        cls = super().__new__(metacls, name, bases, classdict, **kwds)
+
+        # If the enum class has no members, return it as is.
+        if not cls.__members__:
+            return cls
+
+        # Add the indices attribute to the enum class.
+        cls.indices = numpy.arange(len(cls), dtype=t.EnumDType)
+
+        # Add the names attribute to the enum class.
+        cls.names = numpy.array(cls._member_names_, dtype=t.StrDType)
+
+        # Add the enums attribute to the enum class.
+        cls.enums = numpy.array(cls, dtype=t.ObjDType)
+
+        # Return the modified enum class.
+        return cls
+
+    def __dir__(cls) -> list[str]:
+        return sorted({"indices", "names", "enums", *super().__dir__()})
+
+    def __hash__(cls) -> int:
+        return object.__hash__(cls.__name__)
+
+    def __eq__(cls, other: object) -> bool:
+        return hash(cls) == hash(other)
+
+    def __ne__(cls, other: object) -> bool:
+        return hash(cls) != hash(other)
+
+
+__all__ = ["EnumType"]

openfisca_core/indexed_enums/_errors.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from . import types as t
+
+
+class EnumEncodingError(TypeError):
+    """Raised when an enum is encoded with an unsupported type."""
+
+    def __init__(
+        self, enum_class: type[t.Enum], value: t.VarArray | t.ArrayLike[object]
+    ) -> None:
+        msg = (
+            f"Failed to encode \"{value}\" of type '{value[0].__class__.__name__}', "
+            "as it is not supported. Please, try again with an array of "
+            f"'{int.__name__}', '{str.__name__}', or '{enum_class.__name__}'."
+        )
+        super().__init__(msg)
+
+
+class EnumMemberNotFoundError(IndexError):
+    """Raised when a member is not found in an enum."""
+
+    def __init__(self, enum_class: type[t.Enum]) -> None:
+        index = [str(enum.index) for enum in enum_class]
+        names = [enum.name for enum in enum_class]
+        msg = (
+            f"Some members were not found in enum '{enum_class.__name__}'. "
+            f"Possible values are: {', '.join(names[:-1])}, and {names[-1]!s}; "
+            f"or their corresponding indices: {', '.join(index[:-1])}, and "
+            f"{index[-1]}."
+        )
+        super().__init__(msg)
+
+
+__all__ = ["EnumEncodingError", "EnumMemberNotFoundError"]