test: fix doctest of disk storage (#1249)

openfisca_core/data_storage/in_memory_storage.py

@@ -5,7 +5,7 @@
 
 
 class InMemoryStorage:
-    """Low-level class responsible for storing and retrieving calculated vectors in memory."""
+    """Storing and retrieving calculated vectors in memory."""
 
     def __init__(self, is_eternal=False) -> None:
         self._arrays = {}
@@ -14,6 +14,8 @@ def __init__(self, is_eternal=False) -> None:
     def get(self, period):
         """
         Examples:
+            >>> import numpy
+
             >>> from openfisca_core import data_storage, periods
 
             >>> storage = data_storage.InMemoryStorage()
@@ -40,6 +42,8 @@ def get(self, period):
     def put(self, value, period) -> None:
         """
         Examples:
+            >>> import numpy
+
             >>> from openfisca_core import data_storage, periods
 
             >>> storage = data_storage.InMemoryStorage()
@@ -63,6 +67,8 @@ def put(self, value, period) -> None:
     def delete(self, period=None) -> None:
         """
         Examples:
+            >>> import numpy
+
             >>> from openfisca_core import data_storage, periods
 
             >>> storage = data_storage.InMemoryStorage()

openfisca_core/data_storage/on_disk_storage.py

@@ -9,7 +9,7 @@
 
 
 class OnDiskStorage:
-    """Low-level class responsible for storing and retrieving calculated vectors on disk."""
+    """Storing and retrieving calculated vectors on disk."""
 
     def __init__(
         self, storage_dir, is_eternal=False, preserve_storage_dir=False
@@ -21,12 +21,59 @@ def __init__(
         self.storage_dir = storage_dir
 
     def _decode_file(self, file):
+        """
+        Examples
+            >>> import tempfile
+
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, indexed_enums, periods
+
+            >>> class Housing(indexed_enums.Enum):
+            ...     OWNER = "Owner"
+            ...     TENANT = "Tenant"
+            ...     FREE_LODGER = "Free lodger"
+            ...     HOMELESS = "Homeless"
+
+            >>> array = numpy.array([1])
+            >>> value = indexed_enums.EnumArray(array, Housing)
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.put(value, period)
+            ...     storage._decode_file(storage._files[period])
+            EnumArray([<Housing.TENANT: 'Tenant'>])
+
+        """
+
         enum = self._enums.get(file)
         if enum is not None:
             return EnumArray(numpy.load(file), enum)
         return numpy.load(file)
 
     def get(self, period):
+        """
+        Examples:
+            >>> import tempfile
+
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, periods
+
+            >>> value = numpy.array([1, 2, 3])
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.put(value, period)
+            ...     storage.get(period)
+            array([1, 2, 3])
+
+        """
+
         if self.is_eternal:
             period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
@@ -37,6 +84,26 @@ def get(self, period):
         return self._decode_file(values)
 
     def put(self, value, period) -> None:
+        """
+        Examples:
+            >>> import tempfile
+
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, periods
+
+            >>> value = numpy.array([1, "2", "salary"])
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.put(value, period)
+            ...     storage.get(period)
+            array(['1', '2', 'salary'], dtype='<U21')
+
+        """
+
         if self.is_eternal:
             period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
@@ -50,6 +117,38 @@ def put(self, value, period) -> None:
         self._files[period] = path
 
     def delete(self, period=None) -> None:
+        """
+        Examples:
+            >>> import tempfile
+
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, periods
+
+            >>> value = numpy.array([1, 2, 3])
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.put(value, period)
+            ...     storage.get(period)
+            array([1, 2, 3])
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.put(value, period)
+            ...     storage.delete(period)
+            ...     storage.get(period)
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.put(value, period)
+            ...     storage.delete()
+            ...     storage.get(period)
+
+        """
+
         if period is None:
             self._files = {}
             return
@@ -66,9 +165,63 @@ def delete(self, period=None) -> None:
             }
 
     def get_known_periods(self):
+        """
+        Examples:
+            >>> import tempfile
+
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, periods
+
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.get_known_periods()
+            dict_keys([])
+
+            >>> with tempfile.TemporaryDirectory() as directory:
+            ...     storage = data_storage.OnDiskStorage(directory)
+            ...     storage.put([], period)
+            ...     storage.get_known_periods()
+            dict_keys([Period(('year', Instant((2017, 1, 1)), 1))])
+
+        """
+
         return self._files.keys()
 
     def restore(self) -> None:
+        """
+        Examples:
+            >>> import tempfile
+
+            >>> import numpy
+
+            >>> from openfisca_core import data_storage, periods
+
+            >>> value = numpy.array([1, 2, 3])
+            >>> instant = periods.Instant((2017, 1, 1))
+            >>> period = periods.Period(("year", instant, 1))
+            >>> directory = tempfile.TemporaryDirectory()
+
+            >>> storage1 = data_storage.OnDiskStorage(directory.name)
+            >>> storage1.put(value, period)
+            >>> storage1._files
+            {Period(('year', Instant((2017, 1, 1)), 1)): '.../2017.npy'}
+
+            >>> storage2 = data_storage.OnDiskStorage(directory.name)
+            >>> storage2._files
+            {}
+
+            >>> storage2.restore()
+            >>> storage2._files
+            {Period((<DateUnit.YEAR: 'year'>, Instant((2017, 1, 1.../2017.npy'}
+
+            >>> directory.cleanup()
+
+        """
+
         self._files = files = {}
         # Restore self._files from content of storage_dir.
         for filename in os.listdir(self.storage_dir):