Source code for openfisca_core.variables.helpers

from __future__ import annotations

import sortedcontainers
from typing import Optional

from openfisca_core.periods import Period

from .. import variables


[docs]def get_annualized_variable(variable: variables.Variable, annualization_period: Optional[Period] = None) -> variables.Variable: """ Returns a clone of ``variable`` that is annualized for the period ``annualization_period``. When annualized, a variable's formula is only called for a January calculation, and the results for other months are assumed to be identical. """ def make_annual_formula(original_formula, annualization_period = None): def annual_formula(population, period, parameters): if period.start.month != 1 and (annualization_period is None or annualization_period.contains(period)): return population(variable.name, period.this_year.first_month) if original_formula.__code__.co_argcount == 2: return original_formula(population, period) return original_formula(population, period, parameters) return annual_formula new_variable = variable.clone() new_variable.formulas = sortedcontainers.sorteddict.SortedDict({ key: make_annual_formula(formula, annualization_period) for key, formula in variable.formulas.items() }) return new_variable
[docs]def get_neutralized_variable(variable): """ Return a new neutralized variable (to be used by reforms). A neutralized variable always returns its default value, and does not cache anything. """ result = variable.clone() result.is_neutralized = True result.label = '[Neutralized]' if variable.label is None else '[Neutralized] {}'.format(variable.label), return result
def _partition(dict, predicate): true_dict = {} false_dict = {} for key, value in dict.items(): if predicate(key, value): true_dict[key] = value else: false_dict[key] = value return true_dict, false_dict