Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve dump restore #1292

Merged
merged 7 commits into from
Nov 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions .github/workflows/_test-pip.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,9 @@ jobs:
${{ inputs.activate_command }}
make test-country

- name: Run Extension Template tests
if: ${{ startsWith(inputs.os, 'ubuntu') }}
run: |
${{ inputs.activate_command }}
make test-extension
## There are no tests in extension template anymore
eraviart marked this conversation as resolved.
Show resolved Hide resolved
# - name: Run Extension Template tests
# if: ${{ startsWith(inputs.os, 'ubuntu') }}
# run: |
# ${{ inputs.activate_command }}
# make test-extension
6 changes: 6 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,11 @@
# Changelog

### 43.2.3 [#1292](https://github.com/openfisca/openfisca-core/pull/1292)

#### Technical changes

- Improve the way Enum arrays are loaded and saved in dumps

### 43.2.2 [#1280](https://github.com/openfisca/openfisca-core/pull/1280)

#### Documentation
Expand Down
7 changes: 4 additions & 3 deletions openfisca_core/data_storage/on_disk_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,10 @@ def __init__(
storage_dir: str,
is_eternal: bool = False,
preserve_storage_dir: bool = False,
enums: MutableMapping[str, type[t.Enum]] | None = None,
) -> None:
self._files = {}
self._enums = {}
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
Expand Down Expand Up @@ -90,7 +91,7 @@ def _decode_file(self, file: str) -> t.Array[t.DTypeGeneric]:
EnumArray([Housing.TENANT])

"""
enum = self._enums.get(file)
enum = self._enums.get(self.storage_dir)

if enum is not None:
return EnumArray(numpy.load(file), enum)
Expand Down Expand Up @@ -169,7 +170,7 @@ def put(self, value: t.Array[t.DTypeGeneric], period: None | t.Period) -> None:
filename = str(period)
path = os.path.join(self.storage_dir, filename) + ".npy"
if isinstance(value, EnumArray) and value.possible_values is not None:
self._enums[path] = value.possible_values
self._enums[self.storage_dir] = value.possible_values
value = value.view(numpy.ndarray)
numpy.save(path, value)
self._files[period] = path
Expand Down
26 changes: 7 additions & 19 deletions openfisca_core/indexed_enums/enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,28 +117,16 @@ def __init__(self, *__args: object, **__kwargs: object) -> None:
"""
self.index = len(self._member_names_)

# Bypass the slow Enum.__eq__
__eq__ = object.__eq__

# In Python 3, __hash__ must be defined if __eq__ is defined to stay
# hashable.
__hash__ = object.__hash__

def __repr__(self) -> str:
return f"{self.__class__.__name__}.{self.name}"

def __hash__(self) -> int:
return object.__hash__(self.__class__.__name__ + self.name)

def __eq__(self, other: object) -> bool:
if (
isinstance(other, Enum)
and self.__class__.__name__ == other.__class__.__name__
):
return self.index == other.index
return NotImplemented

def __ne__(self, other: object) -> bool:
if (
isinstance(other, Enum)
and self.__class__.__name__ == other.__class__.__name__
):
return self.index != other.index
return NotImplemented

@classmethod
def encode(cls, array: t.VarArray | t.ArrayLike[object]) -> t.EnumArray:
"""Encode an encodable array into an :class:`.EnumArray`.
Expand Down
22 changes: 14 additions & 8 deletions openfisca_core/tools/simulation_dumper.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import numpy

from openfisca_core.data_storage import OnDiskStorage
from openfisca_core.indexed_enums import Enum
from openfisca_core.periods import DateUnit
from openfisca_core.simulations import Simulation

Expand Down Expand Up @@ -117,21 +118,26 @@ def _restore_entity(population, directory):
return person_count


def _restore_holder(simulation, variable, directory) -> None:
storage_dir = os.path.join(directory, variable)
is_variable_eternal = (
simulation.tax_benefit_system.get_variable(variable).definition_period
== DateUnit.ETERNITY
)
def _restore_holder(simulation, variable_name, directory) -> None:
storage_dir = os.path.join(directory, variable_name)

holder = simulation.get_holder(variable_name)

is_variable_eternal = holder.variable.definition_period == DateUnit.ETERNITY

disk_storage = OnDiskStorage(
storage_dir,
is_eternal=is_variable_eternal,
preserve_storage_dir=True,
enums=(
{storage_dir: holder.variable.possible_values}
if holder.variable.value_type == Enum
and holder.variable.possible_values is not None
else {}
),
)
disk_storage.restore()

holder = simulation.get_holder(variable)

for period in disk_storage.get_known_periods():
value = disk_storage.get(period)
holder.put_in_cache(value, period)
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@

setup(
name="OpenFisca-Core",
version="43.2.2",
version="43.2.3",
author="OpenFisca Team",
author_email="contact@openfisca.org",
classifiers=[
Expand Down
3 changes: 3 additions & 0 deletions tests/core/test_yaml.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,9 @@ def test_shell_script_with_reform() -> None:

# TODO(Mauko Quiroga-Alvarado): Fix this test
# https://github.com/openfisca/openfisca-core/issues/962
@pytest.mark.skip(
reason="Does not work because tests are not in openfisca_extension_template anymore."
)
@pytest.mark.skipif(sys.platform == "win32", reason="Does not work on Windows.")
def test_shell_script_with_extension() -> None:
tests_dir = os.path.join(openfisca_extension_template.__path__[0], "tests")
Expand Down
Loading