Fix data storage types

openfisca_core/data_storage/on_disk_storage.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, KeysView, Union
+from typing import Any, AbstractSet, MutableMapping
 from openfisca_core.types import ArrayType
 
 import os
@@ -9,7 +9,7 @@
 import numpy
 
 from openfisca_core import periods
-from openfisca_core.indexed_enums import EnumArray
+from openfisca_core.indexed_enums import Enum, EnumArray
 
 
 class OnDiskStorage:
@@ -24,20 +24,21 @@ def __init__(
             preserve_storage_dir: bool = False,
             ) -> None:
 
-        self._files = {}
-        self._enums = {}
+        self._files: MutableMapping[periods.Period, ArrayType[Any]] = {}
+        self._enums: MutableMapping[str, Enum] = {}
         self.is_eternal = is_eternal
         self.preserve_storage_dir = preserve_storage_dir
         self.storage_dir = storage_dir
 
     def _decode_file(self, file):
+        breakpoint()
         enum = self._enums.get(file)
         if enum is not None:
             return EnumArray(numpy.load(file), enum)
         else:
             return numpy.load(file)
 
-    def get(self, period: Union[str, periods.Period]) -> ArrayType[Any]:
+    def get(self, period: periods.Period) -> ArrayType[Any]:
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
         period = periods.period(period)
@@ -76,10 +77,11 @@ def delete(self, period = None):
                 if not period.contains(period_item)
                 }
 
-    def get_known_periods(self) -> KeysView[str]:
+    def get_known_periods(self) -> AbstractSet[periods.Period]:
         return self._files.keys()
 
     def restore(self) -> None:
+        files: MutableMapping[periods.Period, ArrayType[Any]]
         self._files = files = {}
         # Restore self._files from content of storage_dir.
         for filename in os.listdir(self.storage_dir):