Fix doc & type definitions in the entities module (#1220)

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+### 41.5.5 [#1220](https://github.com/openfisca/openfisca-core/pull/1220)
+
+#### Technical changes
+
+- Fix doc & type definitions in the entities module
+
 ### 41.5.4 [#1219](https://github.com/openfisca/openfisca-core/pull/1219)
 
 #### Technical changes

openfisca_core/entities/__init__.py

@@ -27,4 +27,14 @@
 from .helpers import build_entity, find_role
 from .role import Role
 
-__all__ = ["Entity", "GroupEntity", "Role", "build_entity", "find_role", "types"]
+SingleEntity = Entity
+
+__all__ = [
+    "Entity",
+    "SingleEntity",
+    "GroupEntity",
+    "Role",
+    "build_entity",
+    "find_role",
+    "types",
+]

openfisca_core/entities/_core_entity.py

@@ -1,24 +1,19 @@
 from __future__ import annotations
 
-from typing import Any
-
-from openfisca_core.types import TaxBenefitSystem, Variable
-
 import os
 from abc import abstractmethod
 
+from . import types as t
 from .role import Role
-from .types import Entity
 
 
 class _CoreEntity:
     """Base class to build entities from."""
 
     #: A key to identify the entity.
-    key: str
-
+    key: t.EntityKey
     #: The ``key``, pluralised.
-    plural: str | None
+    plural: t.EntityPlural | None
 
     #: A summary description.
     label: str | None
@@ -30,31 +25,37 @@ class _CoreEntity:
     is_person: bool
 
     #: A TaxBenefitSystem instance.
-    _tax_benefit_system: TaxBenefitSystem | None = None
+    _tax_benefit_system: t.TaxBenefitSystem | None = None
 
     @abstractmethod
-    def __init__(self, key: str, plural: str, label: str, doc: str, *args: Any) -> None:
+    def __init__(
+        self, key: str, plural: str, label: str, doc: str, *args: object
+    ) -> None:
         ...
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.key})"
 
-    def set_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem) -> None:
+    def set_tax_benefit_system(self, tax_benefit_system: t.TaxBenefitSystem) -> None:
         """An Entity belongs to a TaxBenefitSystem."""
         self._tax_benefit_system = tax_benefit_system
 
     def get_variable(
         self,
         variable_name: str,
         check_existence: bool = False,
-    ) -> Variable | None:
+    ) -> t.Variable | None:
         """Get a ``variable_name`` from ``variables``."""
+        if self._tax_benefit_system is None:
+            raise ValueError(
+                "You must set 'tax_benefit_system' before calling this method."
+            )
         return self._tax_benefit_system.get_variable(variable_name, check_existence)
 
     def check_variable_defined_for_entity(self, variable_name: str) -> None:
         """Check if ``variable_name`` is defined for ``self``."""
-        variable: Variable | None
-        entity: Entity
+        variable: t.Variable | None
+        entity: t.CoreEntity
 
         variable = self.get_variable(variable_name, check_existence=True)
 
@@ -71,7 +72,7 @@ def check_variable_defined_for_entity(self, variable_name: str) -> None:
             )
             raise ValueError(os.linesep.join(message))
 
-    def check_role_validity(self, role: Any) -> None:
+    def check_role_validity(self, role: object) -> None:
         """Check if a ``role`` is an instance of Role."""
         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,5 +1,6 @@
 import textwrap
 
+from . import types as t
 from ._core_entity import _CoreEntity
 
 
@@ -9,8 +10,8 @@ class Entity(_CoreEntity):
     """
 
     def __init__(self, key: str, plural: str, label: str, doc: str) -> None:
-        self.key = key
+        self.key = t.EntityKey(key)
         self.label = label
-        self.plural = plural
+        self.plural = t.EntityPlural(plural)
         self.doc = textwrap.dedent(doc)
         self.is_person = True

openfisca_core/entities/group_entity.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Iterable, Mapping
-from typing import Any
+from collections.abc import Iterable, Sequence
 
 import textwrap
 from itertools import chain
 
+from . import types as t
 from ._core_entity import _CoreEntity
 from .role import Role
 
@@ -34,12 +34,12 @@ def __init__(
         plural: str,
         label: str,
         doc: str,
-        roles: Iterable[Mapping[str, Any]],
+        roles: Sequence[t.RoleParams],
         containing_entities: Iterable[str] = (),
     ) -> None:
-        self.key = key
+        self.key = t.EntityKey(key)
         self.label = label
-        self.plural = plural
+        self.plural = t.EntityPlural(plural)
         self.doc = textwrap.dedent(doc)
         self.is_person = False
         self.roles_description = roles

openfisca_core/entities/helpers.py

@@ -1,22 +1,23 @@
 from __future__ import annotations
 
-from typing import Iterable
+from collections.abc import Iterable, Sequence
+from typing import Optional
 
+from . import types as t
 from .entity import Entity
 from .group_entity import GroupEntity
-from .types import Role
 
 
 def build_entity(
-    key,
-    plural,
-    label,
-    doc="",
-    roles=None,
-    is_person=False,
-    class_override=None,
-    containing_entities=(),
-):
+    key: str,
+    plural: str,
+    label: str,
+    doc: str = "",
+    roles: Optional[Sequence[t.RoleParams]] = None,
+    is_person: bool = False,
+    class_override: object | None = None,
+    containing_entities: Sequence[str] = (),
+) -> t.SingleEntity | t.GroupEntity:
     """Build a SingleEntity or a GroupEntity.
 
     Args:
@@ -35,7 +36,7 @@ def build_entity(
         :obj:`.GroupEntity`: When ``is_person`` is False.
 
     Raises:
-        ValueError if ``roles`` is not a sequence.
+        NotImplementedError: if ``roles`` is None.
 
     Examples:
         >>> from openfisca_core import entities
@@ -73,15 +74,17 @@ def build_entity(
     if is_person:
         return Entity(key, plural, label, doc)
 
-    else:
+    if roles is not None:
         return GroupEntity(
             key, plural, label, doc, roles, containing_entities=containing_entities
         )
 
+    raise NotImplementedError
+
 
 def find_role(
-    roles: Iterable[Role], key: str, *, total: int | None = None
-) -> Role | None:
+    roles: Iterable[t.Role], key: t.RoleKey, *, total: int | None = None
+) -> t.Role | None:
     """Find a Role in a GroupEntity.
 
     Args:

openfisca_core/entities/role.py

@@ -6,7 +6,7 @@
 import dataclasses
 import textwrap
 
-from .types import Entity
+from .types import SingleEntity
 
 
 class Role:
@@ -48,7 +48,7 @@ class Role:
     """
 
     #: The Entity the Role belongs to.
-    entity: Entity
+    entity: SingleEntity
 
     #: A description of the Role.
     description: _Description
@@ -79,7 +79,7 @@ def doc(self) -> str | None:
         """A full description, non-indented."""
         return self.description.doc
 
-    def __init__(self, description: Mapping[str, Any], entity: Entity) -> None:
+    def __init__(self, description: Mapping[str, Any], entity: SingleEntity) -> None:
         self.description = _Description(
             **{
                 key: value

openfisca_core/entities/types.py

@@ -1,30 +1,61 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import Protocol, TypedDict
+from typing import NewType, Protocol
+from typing_extensions import Required, TypedDict
 
+from openfisca_core import types as t
 
-class Entity(Protocol):
+# Entities
+
+#: For example "person".
+EntityKey = NewType("EntityKey", str)
+
+#: For example "persons".
+EntityPlural = NewType("EntityPlural", str)
+
+#: For example "principal".
+RoleKey = NewType("RoleKey", str)
+
+#: For example "parents".
+RolePlural = NewType("RolePlural", str)
+
+
+class CoreEntity(t.CoreEntity, Protocol):
+    key: EntityKey
+    plural: EntityPlural | None
+
+
+class SingleEntity(t.SingleEntity, Protocol):
     ...
 
 
-class GroupEntity(Protocol):
+class GroupEntity(t.GroupEntity, Protocol):
     ...
 
 
-class Role(Protocol):
-    max: int | None
+class Role(t.Role, Protocol):
     subroles: Iterable[Role] | None
 
-    @property
-    def key(self) -> str:
-        ...
-
 
 class RoleParams(TypedDict, total=False):
-    key: str
+    key: Required[str]
     plural: str
     label: str
     doc: str
     max: int
     subroles: list[str]
+
+
+# Tax-Benefit systems
+
+
+class TaxBenefitSystem(t.TaxBenefitSystem, Protocol):
+    ...
+
+
+# Variables
+
+
+class Variable(t.Variable, Protocol):
+    ...

openfisca_core/populations/population.py

@@ -3,7 +3,7 @@
 from typing import Dict, NamedTuple, Optional, Sequence, Union
 from typing_extensions import TypedDict
 
-from openfisca_core.types import Array, Entity, Period, Role, Simulation
+from openfisca_core.types import Array, Period, Role, Simulation, SingleEntity
 
 import traceback
 
@@ -16,12 +16,12 @@
 
 class Population:
     simulation: Optional[Simulation]
-    entity: Entity
+    entity: SingleEntity
     _holders: Dict[str, holders.Holder]
     count: int
     ids: Array[str]
 
-    def __init__(self, entity: Entity) -> None:
+    def __init__(self, entity: SingleEntity) -> None:
         self.simulation = None
         self.entity = entity
         self._holders = {}

openfisca_core/projectors/helpers.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Mapping
 
-from openfisca_core.entities.types import Entity, GroupEntity, Role
+from openfisca_core.types import GroupEntity, Role, SingleEntity
 
 from openfisca_core import entities, projectors
 
@@ -110,7 +110,7 @@ def get_projector_from_shortcut(
 
     """
 
-    entity: Entity | GroupEntity = population.entity
+    entity: SingleEntity | GroupEntity = population.entity
 
     if isinstance(entity, entities.Entity):
         populations: Mapping[

openfisca_core/projectors/typing.py

@@ -3,12 +3,12 @@
 from collections.abc import Mapping
 from typing import Protocol
 
-from openfisca_core.entities.types import Entity, GroupEntity
+from openfisca_core.types import GroupEntity, SingleEntity
 
 
 class Population(Protocol):
     @property
-    def entity(self) -> Entity:
+    def entity(self) -> SingleEntity:
         ...
 
     @property