Skip to content

Commit

Permalink
docs(core_population): fix types
Browse files Browse the repository at this point in the history
  • Loading branch information
bonjourmauko committed Nov 21, 2024
1 parent c14cfb9 commit 858d2e3
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 12 deletions.
13 changes: 9 additions & 4 deletions openfisca_core/populations/_core_population.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,12 @@ class CorePopulation:
#: The :class:`.Simulation` for which the population is calculated.
simulation: None | t.Simulation = None

#: The holders of the variables.
_holders: t.HolderByVariable[t.VarDType]

def __init__(self, entity: t.CoreEntity, *__args: object, **__kwds: object) -> None:
self.entity = entity
self._holders: t.HolderByVariable = {}
self._holders = {}

def __call__(
self,
Expand Down Expand Up @@ -337,7 +340,7 @@ def check_period_validity(

# Helpers

def get_holder(self, variable_name: t.VariableName) -> t.Holder:
def get_holder(self, variable_name: t.VariableName) -> t.Holder[t.VarDType]:
"""Return the holder of a variable.
Args:
Expand Down Expand Up @@ -388,8 +391,10 @@ def get_holder(self, variable_name: t.VariableName) -> t.Holder:
if holder:
return holder
variable = self.entity.get_variable(variable_name)
self._holders[variable_name] = holder = holders.Holder(variable, self)
return holder
if variable is None:
raise NotImplementedError
self._holders[variable_name] = holders.Holder(variable, self)
return self._holders[variable_name]

def get_memory_usage(
self,
Expand Down
18 changes: 10 additions & 8 deletions openfisca_core/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,21 +254,21 @@ def _on_disk_storable(self, /) -> bool: ...
def _do_not_store(self, /) -> bool: ...
def clone(self, __population: CorePopulation, /) -> Holder[_N]: ...
def create_disk_storage(
self, /, __dir: None | str = ..., __preserve: bool = ...
self, __dir: None | str = ..., __preserve: bool = ..., /
) -> OnDiskStorage[_N]: ...
def delete_arrays(self, /, __period: None | Period = None) -> None: ...
def get_array(self, /, __period: Period) -> None | Array[_N]: ...
def delete_arrays(self, __period: None | Period = None, /) -> None: ...
def get_array(self, __period: Period, /) -> None | Array[_N]: ...
def get_memory_usage(self, /) -> MemoryUsage: ...
def get_known_periods(self) -> list[Period]: ...
def set_input(
self,
/,
__period: Period,
__array: Array[_N] | ArrayLike[_L],
/,
) -> None | Array[_N]: ...
def put_in_cache(self, /, __value: Array[_N], period: Period) -> None: ...
def put_in_cache(self, __value: Array[_N], period: Period, /) -> None: ...
def default_array(self, /) -> Array[_N]: ...
def _set(self, /, __period: None | Period, __value: Array[_N]) -> None: ...
def _set(self, __period: None | Period, __value: Array[_N], /) -> None: ...
def _to_array(self, __value: Array[_N] | ArrayLike[_L], /) -> Array[_N]: ...


Expand Down Expand Up @@ -427,9 +427,11 @@ def count(self, /) -> int: ...
@property
def entity(self, /) -> CoreEntity: ...
@property
def ids(self, /) -> StrArray: ...
def ids(self, /) -> ArrayLike[str]: ...
@property
def simulation(self, /) -> None | Simulation: ...
@property
def _holders(self, /) -> HolderByVariable[_N]: ...


class SinglePopulation(CorePopulation, Protocol):
Expand Down Expand Up @@ -464,7 +466,7 @@ def tax_benefit_system(self, /) -> TaxBenefitSystem: ...
def trace(self, /) -> bool: ...
@property
def tracer(self, /) -> FullTracer: ...
def calculate(self, __name: VariableName, period__: Period, /) -> VarArray: ...
def calculate(self, __name: VariableName, __period: Period, /) -> VarArray: ...
def calculate_add(self, __name: VariableName, __period: Period, /) -> VarArray: ...
def calculate_divide(
self, __name: VariableName, __period: Period, /
Expand Down

0 comments on commit 858d2e3

Please sign in to comment.