From 19d26df229a132bc53813fe6c0068c80261a34b6 Mon Sep 17 00:00:00 2001
From: Mauko Quiroga <public@mauko.me>
Date: Tue, 29 Nov 2022 17:05:52 +0100
Subject: [PATCH] Fix types

---
 openfisca_core/holders/__init__.py       |  1 +
 openfisca_core/holders/holder.py         |  4 ++--
 openfisca_core/populations/population.py | 10 +---------
 openfisca_core/types/_domain.py          |  9 +++++++++
 4 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/openfisca_core/holders/__init__.py b/openfisca_core/holders/__init__.py
index a7d46e38a6..8bd7722ba3 100644
--- a/openfisca_core/holders/__init__.py
+++ b/openfisca_core/holders/__init__.py
@@ -23,3 +23,4 @@
 
 from .helpers import set_input_dispatch_by_period, set_input_divide_by_period  # noqa: F401
 from .holder import Holder  # noqa: F401
+from .memory_usage import MemoryUsage  # noqa: F401
diff --git a/openfisca_core/holders/holder.py b/openfisca_core/holders/holder.py
index bee8e5c495..ae7e3fbcec 100644
--- a/openfisca_core/holders/holder.py
+++ b/openfisca_core/holders/holder.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Sequence, Union
+from typing import Any, Optional, Sequence, Union
 
 import os
 import warnings
@@ -166,7 +166,7 @@ def set_input(
             self,
             period: types.Period,
             array: Union[numpy.ndarray, Sequence[Any]],
-            ) -> numpy.ndarray:
+            ) -> Optional[numpy.ndarray]:
         """Set a Variable's array of values of a given Period.
 
         Args:
diff --git a/openfisca_core/populations/population.py b/openfisca_core/populations/population.py
index 68d85eb5c7..cb243aff70 100644
--- a/openfisca_core/populations/population.py
+++ b/openfisca_core/populations/population.py
@@ -8,7 +8,7 @@
 import numpy
 
 from openfisca_core import periods, projectors
-from openfisca_core.holders import Holder
+from openfisca_core.holders import Holder, MemoryUsage
 from openfisca_core.projectors import Projector
 from openfisca_core.types import Array, Entity, Period, Role, Simulation
 
@@ -272,11 +272,3 @@ class Calculate(NamedTuple):
 class MemoryUsageByVariable(TypedDict, total = False):
     by_variable: Dict[str, MemoryUsage]
     total_nb_bytes: int
-
-
-class MemoryUsage(TypedDict, total = False):
-    cell_size: int
-    dtype: numpy.dtype
-    nb_arrays: int
-    nb_cells_by_array: int
-    total_nb_bytes: int
diff --git a/openfisca_core/types/_domain.py b/openfisca_core/types/_domain.py
index 3507b42b46..404c133a29 100644
--- a/openfisca_core/types/_domain.py
+++ b/openfisca_core/types/_domain.py
@@ -41,10 +41,19 @@ def get_memory_usage(self) -> Any:
         """Abstract method."""
 
 
+class Instant(Protocol):
+    """Instant protocol."""
+
+
 @typing_extensions.runtime_checkable
 class Period(Protocol):
     """Period protocol."""
 
+    @property
+    @abc.abstractmethod
+    def unit(self) -> Instant:
+        """Abstract method."""
+
 
 class Population(Protocol):
     """Population protocol."""