doc: document get_projector_from_shortcut() (#1197)

CHANGELOG.md

@@ -1,5 +1,15 @@
 # Changelog
 
+## 41.4.0 [#1197](https://github.com/openfisca/openfisca-core/pull/1197)
+
+#### New features
+
+- Add `entities.find_role()` to find roles by key and `max`.
+
+#### Technical changes
+
+- Document `projectors.get_projector_from_shortcut()`.
+
 ## 41.3.0 [#1200](https://github.com/openfisca/openfisca-core/pull/1200)
 
 > As `TracingParameterNodeAtInstant` is a wrapper for `ParameterNodeAtInstant`

openfisca_core/entities/__init__.py

@@ -24,7 +24,7 @@
 from . import typing
 from .entity import Entity
 from .group_entity import GroupEntity
-from .helpers import build_entity
+from .helpers import build_entity, find_role
 from .role import Role
 
-__all__ = ["Entity", "GroupEntity", "Role", "build_entity", "typing"]
+__all__ = ["Entity", "GroupEntity", "Role", "build_entity", "find_role", "typing"]

openfisca_core/entities/_core_entity.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from openfisca_core.types import TaxBenefitSystem, Variable
+from typing import Any
+
+import os
+
+from .role import Role
+from .typing import Entity
+
+
+class _CoreEntity:
+    """Base class to build entities from."""
+
+    #: A key to identify the entity.
+    key: str
+
+    #: The ``key``, pluralised.
+    plural: str | None
+
+    #: A summary description.
+    label: str | None
+
+    #: A full description.
+    doc: str | None
+
+    #: Whether the entity is a person or not.
+    is_person: bool
+
+    #: A TaxBenefitSystem instance.
+    _tax_benefit_system: TaxBenefitSystem | None = None
+
+    @abstractmethod
+    def __init__(self, key: str, plural: str, label: str, doc: str, *args: Any) -> None:
+        ...
+
+    def set_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem) -> None:
+        self._tax_benefit_system = tax_benefit_system
+
+    def get_variable(
+        self,
+        variable_name: str,
+        check_existence: bool = False,
+    ) -> Variable | None:
+        return self._tax_benefit_system.get_variable(variable_name, check_existence)
+
+    def check_variable_defined_for_entity(self, variable_name: str) -> None:
+        variable: Variable | None
+        entity: Entity
+
+        variable = self.get_variable(variable_name, check_existence=True)
+
+        if variable is not None:
+            entity = variable.entity
+
+        if entity.key != self.key:
+            message = (
+                f"You tried to compute the variable '{variable_name}' for",
+                f"the entity '{self.plural}'; however the variable",
+                f"'{variable_name}' is defined for '{entity.plural}'.",
+                "Learn more about entities in our documentation:",
+                "<https://openfisca.org/doc/coding-the-legislation/50_entities.html>.",
+            )
+            raise ValueError(os.linesep.join(message))
+
+    def check_role_validity(self, role: Any) -> None:
+        if role is not None and not isinstance(role, Role):
+            raise ValueError(f"{role} is not a valid role")

openfisca_core/entities/entity.py

@@ -1,13 +1,9 @@
-from openfisca_core.types import TaxBenefitSystem, Variable
-from typing import Any, Optional
-
-import os
 import textwrap
 
-from .role import Role
+from ._core_entity import _CoreEntity
 
 
-class Entity:
+class Entity(_CoreEntity):
     """
     Represents an entity (e.g. a person, a household, etc.) on which calculations can be run.
     """
@@ -18,37 +14,3 @@ def __init__(self, key: str, plural: str, label: str, doc: str) -> None:
         self.plural = plural
         self.doc = textwrap.dedent(doc)
         self.is_person = True
-        self._tax_benefit_system = None
-
-    def set_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem):
-        self._tax_benefit_system = tax_benefit_system
-
-    def check_role_validity(self, role: Any) -> None:
-        if role is not None and not isinstance(role, Role):
-            raise ValueError(f"{role} is not a valid role")
-
-    def get_variable(
-        self,
-        variable_name: str,
-        check_existence: bool = False,
-    ) -> Optional[Variable]:
-        return self._tax_benefit_system.get_variable(variable_name, check_existence)
-
-    def check_variable_defined_for_entity(self, variable_name: str) -> None:
-        variable: Optional[Variable]
-        entity: Entity
-
-        variable = self.get_variable(variable_name, check_existence=True)
-
-        if variable is not None:
-            entity = variable.entity
-
-        if entity.key != self.key:
-            message = (
-                f"You tried to compute the variable '{variable_name}' for",
-                f"the entity '{self.plural}'; however the variable",
-                f"'{variable_name}' is defined for '{entity.plural}'.",
-                "Learn more about entities in our documentation:",
-                "<https://openfisca.org/doc/coding-the-legislation/50_entities.html>.",
-            )
-            raise ValueError(os.linesep.join(message))

openfisca_core/entities/group_entity.py

@@ -1,13 +1,16 @@
+from __future__ import annotations
+
 from collections.abc import Iterable, Mapping
 from typing import Any
 
+import textwrap
 from itertools import chain
 
-from .entity import Entity
+from ._core_entity import _CoreEntity
 from .role import Role
 
 
-class GroupEntity(Entity):
+class GroupEntity(_CoreEntity):
     """Represents an entity containing several others with different roles.
 
     A :class:`.GroupEntity` represents an :class:`.Entity` containing
@@ -34,19 +37,23 @@ def __init__(
         roles: Iterable[Mapping[str, Any]],
         containing_entities: Iterable[str] = (),
     ) -> None:
-        super().__init__(key, plural, label, doc)
+        self.key = key
+        self.label = label
+        self.plural = plural
+        self.doc = textwrap.dedent(doc)
+        self.is_person = False
         self.roles_description = roles
-        self.roles = []
+        self.roles: Iterable[Role] = ()
         for role_description in roles:
             role = Role(role_description, self)
             setattr(self, role.key.upper(), role)
-            self.roles.append(role)
+            self.roles = (*self.roles, role)
             if role_description.get("subroles"):
-                role.subroles = []
+                role.subroles = ()
                 for subrole_key in role_description["subroles"]:
                     subrole = Role({"key": subrole_key, "max": 1}, self)
                     setattr(self, subrole.key.upper(), subrole)
-                    role.subroles.append(subrole)
+                    role.subroles = (*role.subroles, subrole)
                 role.max = len(role.subroles)
         self.flattened_roles = tuple(
             chain.from_iterable(role.subroles or [role] for role in self.roles)

openfisca_core/entities/helpers.py

@@ -1,5 +1,10 @@
+from __future__ import annotations
+
+from typing import Iterable
+
 from .entity import Entity
 from .group_entity import GroupEntity
+from .typing import Role
 
 
 def build_entity(
@@ -18,3 +23,75 @@ def build_entity(
         return GroupEntity(
             key, plural, label, doc, roles, containing_entities=containing_entities
         )
+
+
+def find_role(
+    roles: Iterable[Role], key: str, *, total: int | None = None
+) -> Role | None:
+    """Find a Role in a GroupEntity.
+
+    Args:
+        group_entity (GroupEntity): The entity to search in.
+        key (str): The key of the role to find. Defaults to `None`.
+        total (int | None): The `max` attribute of the role to find.
+
+    Returns:
+        Role | None: The role if found, else `None`.
+
+    Examples:
+        >>> from openfisca_core.entities.typing import RoleParams
+
+        >>> principal = RoleParams(
+        ...     key="principal",
+        ...     label="Principal",
+        ...     doc="Person focus of a calculation in a family context.",
+        ...     max=1,
+        ... )
+
+        >>> partner = RoleParams(
+        ...     key="partner",
+        ...     plural="partners",
+        ...     label="Partners",
+        ...     doc="Persons partners of the principal.",
+        ... )
+
+        >>> parent = RoleParams(
+        ...     key="parent",
+        ...     plural="parents",
+        ...     label="Parents",
+        ...     doc="Persons parents of children of the principal",
+        ...     subroles=["first_parent", "second_parent"],
+        ... )
+
+        >>> group_entity = build_entity(
+        ...     key="family",
+        ...     plural="families",
+        ...     label="Family",
+        ...     doc="A Family represents a collection of related persons.",
+        ...     roles=[principal, partner, parent],
+        ... )
+
+        >>> find_role(group_entity.roles, "principal", total=1)
+        Role(principal)
+
+        >>> find_role(group_entity.roles, "partner")
+        Role(partner)
+
+        >>> find_role(group_entity.roles, "parent", total=2)
+        Role(parent)
+
+        >>> find_role(group_entity.roles, "first_parent", total=1)
+        Role(first_parent)
+
+    """
+
+    for role in roles:
+        if role.subroles:
+            for subrole in role.subroles:
+                if (subrole.max == total) and (subrole.key == key):
+                    return subrole
+
+        if (role.max == total) and (role.key == key):
+            return role
+
+    return None

openfisca_core/entities/role.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping
 from typing import Any
 
 import dataclasses
@@ -57,7 +57,7 @@ class Role:
     max: int | None = None
 
     #: A list of subroles.
-    subroles: Sequence[Role] | None = None
+    subroles: Iterable[Role] | None = None
 
     @property
     def key(self) -> str:

openfisca_core/entities/typing.py

@@ -1,5 +1,30 @@
-from typing import Protocol
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Protocol, TypedDict
 
 
 class Entity(Protocol):
     ...
+
+
+class GroupEntity(Protocol):
+    ...
+
+
+class Role(Protocol):
+    max: int | None
+    subroles: Iterable[Role] | None
+
+    @property
+    def key(self) -> str:
+        ...
+
+
+class RoleParams(TypedDict, total=False):
+    key: str
+    plural: str
+    label: str
+    doc: str
+    max: int
+    subroles: list[str]

openfisca_core/populations/__init__.py

@@ -21,17 +21,27 @@
 #
 # See: https://www.python.org/dev/peps/pep-0008/#imports
 
-from openfisca_core.projectors import (  # noqa: F401
+from openfisca_core.projectors import (
     EntityToPersonProjector,
     FirstPersonToEntityProjector,
     Projector,
     UniqueRoleToEntityProjector,
 )
-from openfisca_core.projectors.helpers import (  # noqa: F401
-    get_projector_from_shortcut,
-    projectable,
-)
+from openfisca_core.projectors.helpers import get_projector_from_shortcut, projectable
+
+from .config import ADD, DIVIDE
+from .group_population import GroupPopulation
+from .population import Population
 
-from .config import ADD, DIVIDE  # noqa: F401
-from .group_population import GroupPopulation  # noqa: F401
-from .population import Population  # noqa: F401
+__all__ = [
+    "ADD",
+    "DIVIDE",
+    "EntityToPersonProjector",
+    "FirstPersonToEntityProjector",
+    "GroupPopulation",
+    "Population",
+    "Projector",
+    "UniqueRoleToEntityProjector",
+    "get_projector_from_shortcut",
+    "projectable",
+]