refactor: make CoreEntity.get_variable reuse .variables

Author: bonjourmauko

openfisca_core/entities/__init__.py

@@ -2,6 +2,7 @@
 
 from . import types
 from ._core_entity import CoreEntity
+from ._errors import TaxBenefitSystemUnsetError, VariableNotFoundError
 from .entity import Entity
 from .group_entity import GroupEntity
 from .helpers import build_entity, find_role
@@ -15,6 +16,8 @@
     "GroupEntity",
     "Role",
     "SingleEntity",
+    "TaxBenefitSystemUnsetError",
+    "VariableNotFoundError",
     "build_entity",
     "find_role",
     "types",

openfisca_core/entities/_core_entity.py

@@ -7,6 +7,7 @@
 import os
 
 from . import types as t
+from ._errors import TaxBenefitSystemUnsetError, VariableNotFoundError
 
 
 class CoreEntity:
@@ -61,7 +62,8 @@ def variables(self, /) -> Mapping[t.VariableName, t.Variable]:
             dict[str, Variable]: The variables defined for the entity.
 
         Raises:
-            ValueError: When the :attr:`.tax_benefit_system` is not set yet.
+            TaxBenefitSystemUnsetError: When the :attr:`.tax_benefit_system` is
+                not set yet.
 
         Examples:
             >>> from openfisca_core import (
@@ -71,12 +73,12 @@ def variables(self, /) -> Mapping[t.VariableName, t.Variable]:
             ...     variables,
             ... )
 
-            >>> this = entities.SingleEntity("this", "", "", "")
-            >>> that = entities.SingleEntity("that", "", "", "")
+            >>> this = entities.SingleEntity("this", "these", "", "")
+            >>> that = entities.SingleEntity("that", "those", "", "")
 
             >>> this.variables
             Traceback (most recent call last):
-            ValueError: You must set 'tax_benefit_system' to call this method.
+            TaxBenefitSystemUnsetError: The tax and benefit system is not se...
 
             >>> tbs = taxbenefitsystems.TaxBenefitSystem([this, that])
             >>> this, that = tbs.entities
@@ -103,8 +105,7 @@ def variables(self, /) -> Mapping[t.VariableName, t.Variable]:
 
         """
         if self.tax_benefit_system is None:
-            msg = "You must set 'tax_benefit_system' to call this method."
-            raise ValueError(msg)
+            raise TaxBenefitSystemUnsetError
         return {
             name: variable
             for name, variable in self.tax_benefit_system.variables.items()
@@ -115,7 +116,7 @@ def get_variable(
         self,
         variable_name: t.VariableName,
         check_existence: bool = False,
-    ) -> t.Variable | None:
+    ) -> None | t.Variable:
         """Get ``variable_name`` from ``variables``.
 
         Args:
@@ -127,9 +128,10 @@ def get_variable(
             None: When the ``Variable`` doesn't exist.
 
         Raises:
-            ValueError: When the :attr:`_tax_benefit_system` is not set yet.
-            ValueError: When ``check_existence`` is ``True`` and
-                the ``Variable`` doesn't exist.
+            TaxBenefitSystemUnsetError: When the :attr:`.tax_benefit_system` is
+                not set yet.
+            VariableNotFoundError: When ``check_existence`` is ``True`` and the
+                ``Variable`` doesn't exist.
 
         Examples:
             >>> from openfisca_core import (
@@ -139,39 +141,44 @@ def get_variable(
             ...     variables,
             ... )
 
-            >>> this = entities.SingleEntity("this", "", "", "")
-            >>> that = entities.SingleEntity("that", "", "", "")
+            >>> this = entities.SingleEntity("this", "these", "", "")
+            >>> that = entities.SingleEntity("that", "those", "", "")
 
             >>> this.get_variable("tax")
             Traceback (most recent call last):
-            ValueError: You must set 'tax_benefit_system' before calling this...
+            TaxBenefitSystemUnsetError: The tax and benefit system is not se...
 
-            >>> tax_benefit_system = taxbenefitsystems.TaxBenefitSystem([this])
-            >>> this.tax_benefit_system = tax_benefit_system
+            >>> tbs = taxbenefitsystems.TaxBenefitSystem([this, that])
+            >>> this, that = tbs.entities
 
             >>> this.get_variable("tax")
 
             >>> this.get_variable("tax", check_existence=True)
             Traceback (most recent call last):
-            VariableNotFoundError: You tried to calculate or to set a value...
+            VariableNotFoundError: You requested the variable 'tax', but it ...
 
             >>> class tax(variables.Variable):
             ...     definition_period = periods.MONTH
             ...     value_type = float
             ...     entity = that
 
-            >>> this.tax_benefit_system.add_variable(tax)
+            >>> tbs.add_variable(tax)
             <openfisca_core.entities._core_entity.tax object at ...>
 
             >>> this.get_variable("tax")
-            <openfisca_core.entities._core_entity.tax object at ...>
 
+            >>> that.get_variable("tax")
+            <openfisca_core.entities._core_entity.tax object at ...>
 
         """
         if self.tax_benefit_system is None:
-            msg = "You must set 'tax_benefit_system' to call this method."
-            raise ValueError(msg)
-        return self.tax_benefit_system.get_variable(variable_name, check_existence)
+            raise TaxBenefitSystemUnsetError
+        if not check_existence:
+            return self.variables.get(variable_name)
+        try:
+            return self.variables[variable_name]
+        except KeyError as error:
+            raise VariableNotFoundError(variable_name, self.plural) from error
 
     def check_variable_defined_for_entity(self, variable_name: t.VariableName) -> None:
         """Check if ``variable_name`` is defined for ``self``.

openfisca_core/entities/_errors.py

@@ -0,0 +1,32 @@
+from . import types as t
+
+
+class TaxBenefitSystemUnsetError(ValueError):
+    """Raised when a tax and benefit system is not set yet."""
+
+    def __init__(self) -> None:
+        msg = (
+            "The tax and benefit system is not set yet. You need to set it "
+            "before using this method: `entity.tax_benefit_system = ...`."
+        )
+        super().__init__(msg)
+
+
+class VariableNotFoundError(IndexError):
+    """Raised when a requested variable is not defined for as entity."""
+
+    def __init__(self, name: t.VariableName, plural: t.EntityPlural) -> None:
+        msg = (
+            f"You requested the variable '{name}', but it was not found in "
+            f"the entity '{plural}'. Are you sure you spelled '{name}' "
+            "correctly? If this code used to work and suddenly does not, "
+            "this is most probably linked to an update of the tax and "
+            "benefit system. Look at its CHANGELOG to learn about renames "
+            "and removals and update your code accordingly."
+            "Learn more about entities in our documentation:",
+            "<https://openfisca.org/doc/coding-the-legislation/50_entities.html>.",
+        )
+        super().__init__(msg)
+
+
+__all__ = ["TaxBenefitSystemUnsetError", "VariableNotFoundError"]