Source code for openfisca_core.variables.variable

from __future__ import annotations

from openfisca_core.types import Formula, Instant
from typing import Optional, Union

import datetime
import re
import textwrap

import numpy
import sortedcontainers

from openfisca_core import periods, tools
from openfisca_core.entities import Entity, GroupEntity
from openfisca_core.indexed_enums import Enum, EnumArray
from openfisca_core.periods import DateUnit, Period

from . import config, helpers

[docs] class Variable: """ A `variable <>`_ of the legislation. Main attributes: .. attribute:: name Name of the variable .. attribute:: value_type The value type of the variable. Possible value types in OpenFisca are ``int`` ``float`` ``bool`` ``str`` ``date`` and ``Enum``. .. attribute:: entity `Entity <,_entities,_role.html>`_ the variable is defined for. For instance : ``Person``, ``Household``. .. attribute:: definition_period `Period <>`_ the variable is defined for. Possible value: ``DateUnit.DAY``, ``DateUnit.MONTH``, ``DateUnit.YEAR``, ``DateUnit.ETERNITY``. .. attribute:: formulas Formulas used to calculate the variable .. attribute:: label Description of the variable .. attribute:: reference Legislative reference describing the variable. .. attribute:: default_value `Default value <>`_ of the variable. Secondary attributes: .. attribute:: baseline_variable If the variable has been introduced in a `reform <>`_ to replace another variable, baseline_variable is the replaced variable. .. attribute:: dtype Numpy `dtype <>`_ used under the hood for the variable. .. attribute:: end `Date <>`_ when the variable disappears from the legislation. .. attribute:: is_neutralized True if the variable is neutralized. Neutralized variables never use their formula, and only return their default values when calculated. .. attribute:: json_type JSON type corresponding to the variable. .. attribute:: max_length If the value type of the variable is ``str``, max length of the string allowed. ``None`` if there is no limit. .. attribute:: possible_values If the value type of the variable is ``Enum``, contains the values the variable can take. .. attribute:: set_input Function used to automatically process variable inputs defined for periods not matching the definition_period of the variable. See more on the `documentation <>`_. Possible values are ``set_input_dispatch_by_period``, ``set_input_divide_by_period``, or nothing. .. attribute:: unit Free text field describing the unit of the variable. Only used as metadata. .. attribute:: documentation Free multilines text field describing the variable context and usage. """ __name__: str def __init__(self, baseline_variable=None): = self.__class__.__name__ attr = { name: value for name, value in self.__class__.__dict__.items() if not name.startswith("__") } self.baseline_variable = baseline_variable self.value_type = self.set( attr, "value_type", required=True, allowed_values=config.VALUE_TYPES.keys() ) self.dtype = config.VALUE_TYPES[self.value_type]["dtype"] self.json_type = config.VALUE_TYPES[self.value_type]["json_type"] if self.value_type == Enum: self.possible_values = self.set( attr, "possible_values", required=True, setter=self.set_possible_values ) if self.value_type == str: self.max_length = self.set(attr, "max_length", allowed_type=int) if self.max_length: self.dtype = "|S{}".format(self.max_length) if self.value_type == Enum: self.default_value = self.set( attr, "default_value", allowed_type=self.possible_values, required=True ) else: self.default_value = self.set( attr, "default_value", allowed_type=self.value_type, default=config.VALUE_TYPES[self.value_type].get("default"), ) self.entity = self.set(attr, "entity", required=True, setter=self.set_entity) self.definition_period = self.set( attr, "definition_period", required=True, allowed_values=DateUnit ) self.label = self.set(attr, "label", allowed_type=str, setter=self.set_label) self.end = self.set(attr, "end", allowed_type=str, setter=self.set_end) self.reference = self.set(attr, "reference", setter=self.set_reference) self.cerfa_field = self.set(attr, "cerfa_field", allowed_type=(str, dict)) self.unit = self.set(attr, "unit", allowed_type=str) self.documentation = self.set( attr, "documentation", allowed_type=str, setter=self.set_documentation ) self.set_input = self.set_set_input(attr.pop("set_input", None)) self.calculate_output = self.set_calculate_output( attr.pop("calculate_output", None) ) self.is_period_size_independent = self.set( attr, "is_period_size_independent", allowed_type=bool, default=config.VALUE_TYPES[self.value_type]["is_period_size_independent"], ) self.introspection_data = self.set( attr, "introspection_data", ) formulas_attr, unexpected_attrs = helpers._partition( attr, lambda name, value: name.startswith(config.FORMULA_NAME_PREFIX) ) self.formulas = self.set_formulas(formulas_attr) if unexpected_attrs: raise ValueError( 'Unexpected attributes in definition of variable "{}": {!r}'.format(, ", ".join(sorted(unexpected_attrs.keys())) ) ) self.is_neutralized = False # ----- Setters used to build the variable ----- # def set( self, attributes, attribute_name, required=False, allowed_values=None, allowed_type=None, setter=None, default=None, ): value = attributes.pop(attribute_name, None) if value is None and self.baseline_variable: return getattr(self.baseline_variable, attribute_name) if required and value is None: raise ValueError( "Missing attribute '{}' in definition of variable '{}'.".format( attribute_name, ) ) if allowed_values is not None and value not in allowed_values: raise ValueError( "Invalid value '{}' for attribute '{}' in variable '{}'. Allowed values are '{}'.".format( value, attribute_name,, allowed_values ) ) if ( allowed_type is not None and value is not None and not isinstance(value, allowed_type) ): if allowed_type == float and isinstance(value, int): value = float(value) else: raise ValueError( "Invalid value '{}' for attribute '{}' in variable '{}'. Must be of type '{}'.".format( value, attribute_name,, allowed_type ) ) if setter is not None: value = setter(value) if value is None and default is not None: return default return value def set_entity(self, entity): if not isinstance(entity, (Entity, GroupEntity)): raise ValueError( f"Invalid value '{entity}' for attribute 'entity' in variable " f"'{}'. Must be an instance of Entity or GroupEntity." ) return entity def set_possible_values(self, possible_values): if not issubclass(possible_values, Enum): raise ValueError( "Invalid value '{}' for attribute 'possible_values' in variable '{}'. Must be a subclass of {}.".format( possible_values,, Enum ) ) return possible_values def set_label(self, label): if label: return label def set_end(self, end): if end: try: return datetime.datetime.strptime(end, "%Y-%m-%d").date() except ValueError: raise ValueError( "Incorrect 'end' attribute format in '{}'. 'YYYY-MM-DD' expected where YYYY, MM and DD are year, month and day. Found: {}".format(, end ) ) def set_reference(self, reference): if reference: if isinstance(reference, str): reference = [reference] elif isinstance(reference, list): pass elif isinstance(reference, tuple): reference = list(reference) else: raise TypeError( "The reference of the variable {} is a {} instead of a String or a List of Strings.".format(, type(reference) ) ) for element in reference: if not isinstance(element, str): raise TypeError( "The reference of the variable {} is a {} instead of a String or a List of Strings.".format(, type(reference) ) ) return reference def set_documentation(self, documentation): if documentation: return textwrap.dedent(documentation) def set_set_input(self, set_input): if not set_input and self.baseline_variable: return self.baseline_variable.set_input return set_input def set_calculate_output(self, calculate_output): if not calculate_output and self.baseline_variable: return self.baseline_variable.calculate_output return calculate_output def set_formulas(self, formulas_attr): formulas = sortedcontainers.sorteddict.SortedDict() for formula_name, formula in formulas_attr.items(): starting_date = self.parse_formula_name(formula_name) if self.end is not None and starting_date > self.end: raise ValueError( 'You declared that "{}" ends on "{}", but you wrote a formula to calculate it from "{}" ({}). The "end" attribute of a variable must be posterior to the start dates of all its formulas.'.format(, self.end, starting_date, formula_name ) ) formulas[str(starting_date)] = formula # If the variable is reforming a baseline variable, keep the formulas from the latter when they are not overridden by new formulas. if self.baseline_variable is not None: first_reform_formula_date = formulas.peekitem(0)[0] if formulas else None formulas.update( { baseline_start_date: baseline_formula for baseline_start_date, baseline_formula in self.baseline_variable.formulas.items() if first_reform_formula_date is None or baseline_start_date < first_reform_formula_date } ) return formulas
[docs] def parse_formula_name(self, attribute_name): """ Returns the starting date of a formula based on its name. Valid dated name formats are : 'formula', 'formula_YYYY', 'formula_YYYY_MM' and 'formula_YYYY_MM_DD' where YYYY, MM and DD are a year, month and day. By convention, the starting date of: - `formula` is `0001-01-01` (minimal date in Python) - `formula_YYYY` is `YYYY-01-01` - `formula_YYYY_MM` is `YYYY-MM-01` """ def raise_error(): raise ValueError( 'Unrecognized formula name in variable "{}". Expecting "formula_YYYY" or "formula_YYYY_MM" or "formula_YYYY_MM_DD where YYYY, MM and DD are year, month and day. Found: "{}".'.format(, attribute_name ) ) if attribute_name == config.FORMULA_NAME_PREFIX: return FORMULA_REGEX = r"formula_(\d{4})(?:_(\d{2}))?(?:_(\d{2}))?$" # YYYY or YYYY_MM or YYYY_MM_DD match = re.match(FORMULA_REGEX, attribute_name) if not match: raise_error() date_str = "-".join( [, or "01", or "01"] ) try: return datetime.datetime.strptime(date_str, "%Y-%m-%d").date() except ValueError: # formula_2005_99_99 for instance raise_error()
# ----- Methods ----- #
[docs] def is_input_variable(self): """ Returns True if the variable is an input variable. """ return len(self.formulas) == 0
@classmethod def get_introspection_data(cls): try: return cls.introspection_data except AttributeError: return "", None, 0
[docs] def get_formula( self, period: Union[Instant, Period, str, int] = None, ) -> Optional[Formula]: """Returns the formula to compute the variable at the given period. If no period is given and the variable has several formulas, the method returns the oldest formula. Args: period: The period to get the formula. Returns: Formula used to compute the variable. """ instant: Optional[Instant] if not self.formulas: return None if period is None: return self.formulas.peekitem(index=0)[ 1 ] # peekitem gets the 1st key-value tuple (the oldest start_date and formula). Return the formula. if isinstance(period, Period): instant = period.start else: try: instant = periods.period(period).start except ValueError: instant = periods.instant(period) if instant is None: return None if self.end and > self.end: return None instant_str = str(instant) for start_date in reversed(self.formulas): if start_date <= instant_str: return self.formulas[start_date] return None
def clone(self): clone = self.__class__() return clone def check_set_value(self, value): if self.value_type == Enum and isinstance(value, str): try: value = self.possible_values[value].index except KeyError: possible_values = [ for item in self.possible_values] raise ValueError( "'{}' is not a known value for '{}'. Possible values are ['{}'].".format( value,, "', '".join(possible_values) ) ) if self.value_type in (float, int) and isinstance(value, str): try: value = tools.eval_expression(value) except SyntaxError: raise ValueError( "I couldn't understand '{}' as a value for '{}'".format( value, ) ) try: value = numpy.array([value], dtype=self.dtype)[0] except (TypeError, ValueError): if self.value_type == error_message = "Can't deal with date: '{}'.".format(value) else: error_message = ( "Can't deal with value: expected type {}, received '{}'.".format( self.json_type, value ) ) raise ValueError(error_message) except OverflowError: error_message = ( "Can't deal with value: '{}', it's too large for type '{}'.".format( value, self.json_type ) ) raise ValueError(error_message) return value def default_array(self, array_size): array = numpy.empty(array_size, dtype=self.dtype) if self.value_type == Enum: array.fill(self.default_value.index) return EnumArray(array, self.possible_values) array.fill(self.default_value) return array