Improve dump restore (#1292)

.github/workflows/_test-pip.yaml

@@ -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
+    # - name: Run Extension Template tests
+    #   if: ${{ startsWith(inputs.os, 'ubuntu') }}
+    #   run: |
+    #     ${{ inputs.activate_command }}
+    #     make test-extension

CHANGELOG.md

@@ -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

openfisca_core/data_storage/on_disk_storage.py

@@ -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
@@ -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)
@@ -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

openfisca_core/indexed_enums/enum.py

@@ -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`.

openfisca_core/tools/simulation_dumper.py

@@ -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
 
@@ -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)

setup.py

@@ -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=[

tests/core/test_yaml.py

@@ -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")