Fix various vectorial parameter access and add unit tests

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+### 43.4.2 [#1336](https://github.com/openfisca/openfisca-core/pull/1336)
+
+#### Technical changes
+
+- Fix various vectorial parameter access and add unit tests
+
 ### 43.4.1 [#1337](https://github.com/openfisca/openfisca-core/pull/1337)
 
 #### Technical changes

openfisca_core/parameters/vectorial_parameter_node_at_instant.py

@@ -14,6 +14,8 @@ def __init__(self, vector, name, instant_str) -> None:
 
  @staticmethod
  def _get_appropriate_subnode_key(node, k):
+ if isinstance(k, numpy.integer):
+ return str(k)
  if isinstance(k, str):
  return k
  if isinstance(k, Enum):
@@ -35,6 +37,8 @@ def _get_appropriate_keys(key):
  [key == item.index for item in enum],
  [item.name for item in enum],
  )
+ if key.dtype == object and issubclass(type(key[0]), Enum):
+ return key
  return key
 
  @staticmethod
@@ -59,9 +63,13 @@ def __getitem__(self, key):
  if isinstance(key, str):
  return getattr(self, key)
 
- nodes = [
- v[a if isinstance(a, str) else a.value] for (v, a) in zip(self.vector, key)
- ]
+ keys = key
+ if isinstance(key[0], Enum):
+ keys = [k.name for k in keys]
+ elif isinstance(key, EnumArray):
+ keys = [key.possible_values.names[v] for v in key]
+
+ nodes = [v[a] for (v, a) in zip(self.vector, keys)]
  return VectorialParameterNodeAtInstant.build_from_nodes(self, nodes)
 
  @staticmethod

setup.py

@@ -70,7 +70,7 @@
 
 setup(
  name="OpenFisca-Core",
- version="43.4.1",
+ version="43.4.2",
  author="OpenFisca Team",
  author_email="contact@openfisca.org",
  classifiers=[

tests/core/parameters_fancy_indexing/multiple_enum_levels.yaml

@@ -0,0 +1,26 @@
+level1:
+ step1:
+ values:
+ 2015-01-01:
+ value: step2
+ step2:
+ values:
+ 2015-01-01:
+ value: step3
+ step3:
+ values:
+ 2015-01-01:
+ value: step3
+level2:
+ step1:
+ values:
+ 2015-01-01:
+ value: step2
+ step2:
+ values:
+ 2015-01-01:
+ value: step3
+ step3:
+ values:
+ 2015-01-01:
+ value: step3

tests/core/parameters_fancy_indexing/test_fancy_indexing.py

@@ -154,6 +154,11 @@ def test_with_properties_starting_by_number() -> None:
  assert_near(P_2[city_code], [100, 300, 200])
 
 
+def test_with_integer() -> None:
+ city_code = numpy.asarray([75012, 75007, 75015])
+ assert_near(P_2[city_code], [100, 300, 200])
+
+
 P_3 = parameters.bareme("2015-01-01")
 
 
@@ -191,9 +196,44 @@ class TypesZone(Enum):
  assert_near(P.single.owner[zone_array], [100, 200, 200, 100])
 
 
+def test_with_enum_in_object_array() -> None:
+ class TypesZone(Enum):
+ z1 = "Zone 1"
+ z2 = "Zone 2"
+
+ zone = numpy.asarray([TypesZone.z1, TypesZone.z2, TypesZone.z2, TypesZone.z1])
+ value = P.single.owner[zone]
+ assert_near(value, [100, 200, 200, 100])
+ # Ensure correct behavior of operator "*"
+ assert_near(value * [2, 1, 1, 2], [200, 200, 200, 200])
+
+
 P_4 = parameters.next_category("2015-01-01")
 
 
 def test_on_leaf_str() -> None:
  current_category = numpy.asarray(["step1", "step1", "step2", "step3"])
  assert_near(P_4[current_category], ["step2", "step2", "step3", "step3"])
+
+
+P_5 = parameters.multiple_enum_levels("2015-01-01")
+
+
+def test_multiple_enums() -> None:
+ class Level(Enum):
+ level1 = "Level 1"
+ level2 = "Level 2"
+
+ class Step(Enum):
+ step1 = "Step 1"
+ step2 = "Step 2"
+ step3 = "Step 3"
+
+ levels = [Level.level1, Level.level1, Level.level2, Level.level2]
+ current_level = EnumArray(numpy.asarray([i.index for i in levels]), Level)
+
+ P_steps = P_5[current_level]
+ steps = [Step.step1, Step.step1, Step.step2, Step.step3]
+ current_step = EnumArray(numpy.asarray([i.index for i in steps]), Step)
+
+ assert_near(P_steps[current_step], ["step2", "step2", "step3", "step3"])