Skip to content

Commit

Permalink
doc(types): add projector types
Browse files Browse the repository at this point in the history
  • Loading branch information
bonjourmauko committed Sep 19, 2024
1 parent 4e3dfcb commit 584e97f
Show file tree
Hide file tree
Showing 5 changed files with 20 additions and 29 deletions.
2 changes: 1 addition & 1 deletion openfisca_core/projectors/entity_to_person_projector.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
class EntityToPersonProjector(Projector):
"""For instance person.family."""

def __init__(self, entity, parent=None):
def __init__(self, entity, parent=None) -> None:
self.reference_entity = entity
self.parent = parent

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
class FirstPersonToEntityProjector(Projector):
"""For instance famille.first_person."""

def __init__(self, entity, parent=None):
def __init__(self, entity, parent=None) -> None:
self.target_entity = entity
self.reference_entity = entity.members
self.parent = parent
Expand Down
18 changes: 8 additions & 10 deletions openfisca_core/projectors/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from openfisca_core import entities, projectors

from .types import GroupEntity, GroupPopulation, Role, SingleEntity, SinglePopulation
from . import types as t


def projectable(function):
Expand All @@ -17,10 +17,10 @@ def projectable(function):


def get_projector_from_shortcut(
population: SinglePopulation | GroupPopulation,
population: t.CorePopulation,
shortcut: str,
parent: projectors.Projector | None = None,
) -> projectors.Projector | None:
parent: t.Projector | None = None,
) -> t.Projector | None:
"""Get a projector from a shortcut.
Projectors are used to project an invidividual Population's or a
Expand Down Expand Up @@ -108,12 +108,10 @@ def get_projector_from_shortcut(
"""

entity: SingleEntity | GroupEntity = population.entity
entity: t.CoreEntity = population.entity

if isinstance(entity, entities.Entity):
populations: Mapping[
str, SinglePopulation | GroupPopulation
] = population.simulation.populations
populations: Mapping[str, t.CorePopulation] = population.simulation.populations

if shortcut not in populations.keys():
return None
Expand All @@ -124,8 +122,8 @@ def get_projector_from_shortcut(
return projectors.FirstPersonToEntityProjector(population, parent)

if isinstance(entity, entities.GroupEntity):
roles: Iterable[Role] = entity.roles
role: Role | None = entities.find_role(roles, shortcut, total=1)
roles: Iterable[t.Role] = entity.roles
role: t.Role | None = entities.find_role(roles, shortcut, total=1)

if role is not None:
return projectors.UniqueRoleToEntityProjector(population, role, parent)
Expand Down
25 changes: 9 additions & 16 deletions openfisca_core/projectors/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,34 +8,27 @@
# Entities


class SingleEntity(t.SingleEntity, Protocol):
...


class GroupEntity(t.GroupEntity, Protocol):
class CoreEntity(t.CoreEntity, Protocol):
...


class Role(t.Role, Protocol):
...


# Populations
# Projectors


class SinglePopulation(t.SinglePopulation, Protocol):
@property
def entity(self) -> t.SingleEntity:
...
class Projector(Protocol):
...

@property
def simulation(self) -> Simulation:
...

# Populations


class GroupPopulation(t.GroupPopulation, Protocol):
class CorePopulation(t.CorePopulation, Protocol):
@property
def entity(self) -> t.GroupEntity:
def entity(self) -> CoreEntity:
...

@property
Expand All @@ -48,5 +41,5 @@ def simulation(self) -> Simulation:

class Simulation(t.Simulation, Protocol):
@property
def populations(self) -> Mapping[str, SinglePopulation | GroupPopulation]:
def populations(self) -> Mapping[str, CorePopulation]:
...
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
class UniqueRoleToEntityProjector(Projector):
"""For instance famille.declarant_principal."""

def __init__(self, entity, role, parent=None):
def __init__(self, entity, role, parent=None) -> None:
self.target_entity = entity
self.reference_entity = entity.members
self.parent = parent
Expand Down

0 comments on commit 584e97f

Please sign in to comment.