Skip to content

Commit

Permalink
feat: add attribute min_value, max_value in variable.py
Browse files Browse the repository at this point in the history
  • Loading branch information
SylviaDu99 committed Jun 3, 2024
1 parent 1ab99f7 commit a6b1a8b
Show file tree
Hide file tree
Showing 3 changed files with 114 additions and 0 deletions.
4 changes: 4 additions & 0 deletions changelog_entry.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
- bump: minor
changes:
added:
- max_value and min_value in Variable class.
31 changes: 31 additions & 0 deletions policyengine_core/variables/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,12 @@ class Variable:
exhaustive_parameter_dependencies: List[str] = None
"""If these parameters (plus the dataset, branch and period) haven't changed, Core will use caching on this variable."""

min_value: (float, int) = None
"""Minimum value of the variable."""

max_value: (float, int) = None
"""Maximum value of the variable."""

def __init__(self, baseline_variable=None):
self.name = self.__class__.__name__
attr = {
Expand Down Expand Up @@ -305,6 +311,19 @@ def __init__(self, baseline_variable=None):
self.exhaustive_parameter_dependencies
]

self.min_value = self.set(
attr,
"min_value",
allowed_type=(float, int),
setter=self.set_min_value,
)
self.max_value = self.set(
attr,
"max_value",
allowed_type=(float, int),
setter=self.set_max_value,
)

formulas_attr, unexpected_attrs = helpers._partition(
attr,
lambda name, value: name.startswith(config.FORMULA_NAME_PREFIX),
Expand Down Expand Up @@ -463,6 +482,18 @@ def set_defined_for(self, defined_for):
defined_for = defined_for.value
return defined_for

def set_min_value(self, min_value):
if min_value is not None:
if self.max_value is not None and min_value > self.max_value:
raise ValueError("min_value cannot be greater than max_value")
return min_value

def set_max_value(self, max_value):
if max_value is not None:
if self.min_value is not None and max_value < self.min_value:
raise ValueError("max_value cannot be smaller than min_value")
return max_value

def parse_formula_name(self, attribute_name):
"""
Returns the starting date of a formula based on its name.
Expand Down
79 changes: 79 additions & 0 deletions tests/core/variables/test_bounds.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
import policyengine_core.country_template as country_template
from policyengine_core.country_template.entities import Person
from policyengine_core.variables import Variable
from policyengine_core.periods import YEAR
from test_variables import get_message

tax_benefit_system = country_template.CountryTaxBenefitSystem()


class variable__bound(Variable):
value_type = int
entity = Person
definition_period = YEAR
label = "Variable with bound."
min_value = 0
max_value = 100


tax_benefit_system.add_variable(variable__bound)


def test_variable__bound():
variable = tax_benefit_system.variables["variable__bound"]
assert variable.min_value == 0
assert variable.max_value == 100


class variable__no_bound(Variable):
value_type = int
entity = Person
definition_period = YEAR
label = "Variable with no bound."


tax_benefit_system.add_variable(variable__no_bound)


def test_variable__no_bound():
variable = tax_benefit_system.variables["variable__no_bound"]
assert variable.min_value is None
assert variable.max_value is None


class variable__small_max(Variable):
value_type = int
entity = Person
definition_period = YEAR
label = "Variable with max value smaller than min value."
min_value = 100
max_value = 99


def test_variable__small_max():
try:
tax_benefit_system.add_variable(variable__small_max)
except ValueError as e:
message = get_message(e)
assert message.startswith("min_value cannot be greater than max_value")
assert not tax_benefit_system.variables.get("variable__small_max")


class variable__wrong_type(Variable):
value_type = int
entity = Person
definition_period = YEAR
label = "Variable with wrong max value data type."
max_value = "string"


def test_variable__wrong_type():
try:
tax_benefit_system.add_variable(variable__wrong_type)
except ValueError as e:
message = get_message(e)
assert message.startswith(
"Invalid value 'string' for attribute 'max_value' in variable 'variable__wrong_type'. "
"Must be of type '(<class 'float'>, <class 'int'>)'."
)
assert not tax_benefit_system.variables.get("variable__wrong_type")

0 comments on commit a6b1a8b

Please sign in to comment.