From bb5ee4229f9f87adbde5e69ab76ab8983ca7de9f Mon Sep 17 00:00:00 2001 From: tmadlener Date: Thu, 20 Jul 2023 15:16:28 +0200 Subject: [PATCH] Add an as_type argument to put_parameter to make explicit typing possible --- python/podio/frame.py | 34 ++++++++++++++++++++++++++-------- python/podio/test_Frame.py | 8 ++++++++ python/podio/test_utils.py | 2 ++ tests/read_python_frame.h | 13 +++++++++++++ 4 files changed, 49 insertions(+), 8 deletions(-) diff --git a/python/podio/frame.py b/python/podio/frame.py index c0c517dfa..925a6a46b 100644 --- a/python/podio/frame.py +++ b/python/podio/frame.py @@ -48,13 +48,19 @@ def _determine_cpp_type(idx_and_type): SUPPORTED_PARAMETER_TYPES = _determine_supported_parameter_types() -def _get_cpp_vector_types(type_str): - """Get the possible std::vector from the passed py_type string.""" - # Gather a list of all types that match the type_str (c++ or python) +def _get_cpp_types(type_str): + """Get all possible c++ types from the passed py_type string.""" types = list(filter(lambda t: type_str in t, SUPPORTED_PARAMETER_TYPES)) if not types: raise ValueError(f'{type_str} cannot be mapped to a valid parameter type') + return types + + +def _get_cpp_vector_types(type_str): + """Get the possible std::vector from the passed py_type string.""" + # Gather a list of all types that match the type_str (c++ or python) + types = _get_cpp_types(type_str) return [f'std::vector<{t}>' for t in map(lambda x: x[0], types)] @@ -196,18 +202,23 @@ def _get_param_value(par_type, name): return _get_param_value(vec_types[0], name) - def put_parameter(self, key, value): + def put_parameter(self, key, value, as_type=None): """Put a parameter into the Frame. Puts a parameter into the Frame after doing some (incomplete) type checks. If a list is passed the parameter type is determined from looking at the first element of the list only. Additionally, since python doesn't differentiate between floats and doubles, floats will always be stored as - double. + doubles by default, use the as_type argument to change this if necessary. Args: key (str): The name of the parameter value (int, float, str or list of these): The parameter value + as_type (str, optional): Explicitly specify the type that should be used + to put the parameter into the Frame. Python types (e.g. "str") will + be converted to c++ types. This will override any automatic type + deduction that happens otherwise. Note that this will be taken at + pretty much face-value and there are only limited checks for this. Raises: ValueError: If a non-supported parameter type is passed @@ -215,9 +226,10 @@ def put_parameter(self, key, value): # For lists we determine the c++ vector type and use that to call the # correct template overload explicitly if isinstance(value, (list, tuple)): - vec_types = _get_cpp_vector_types(type(value[0]).__name__) + type_name = as_type or type(value[0]).__name__ + vec_types = _get_cpp_vector_types(type_name) if len(vec_types) == 0: - raise ValueError(f"Cannot put a parameter of type {type(value[0])} into a Frame") + raise ValueError(f"Cannot put a parameter of type {type_name} into a Frame") par_type = vec_types[0] if isinstance(value[0], float): @@ -226,10 +238,16 @@ def put_parameter(self, key, value): self._frame.putParameter[par_type](key, value) else: + if as_type is not None: + cpp_types = _get_cpp_types(as_type) + if len(cpp_types) == 0: + raise ValueError(f"Cannot put a parameter of type {as_type} into a Frame") + self._frame.putParameter[cpp_types[0]](key, value) + # If we have a single integer, a std::string overload kicks in with higher # priority than the template for some reason. So we explicitly select the # correct template here - if isinstance(value, int): + elif isinstance(value, int): self._frame.putParameter["int"](key, value) else: self._frame.putParameter(key, value) diff --git a/python/podio/test_Frame.py b/python/podio/test_Frame.py index f69c074e9..3d50073c6 100644 --- a/python/podio/test_Frame.py +++ b/python/podio/test_Frame.py @@ -84,6 +84,14 @@ def test_frame_put_parameters(self): self.assertEqual(len(vec), 3) self.assertEqual(vec, [1.23, 4.56, 7.89]) + frame.put_parameter("real_float_vec", [1.23, 4.56, 7.89], as_type="float") + f_vec = frame.get_parameter("real_float_vec", as_type="float") + self.assertEqual(len(f_vec), 3) + self.assertEqual(vec, [1.23, 4.56, 7.89]) + + frame.put_parameter("float_as_float", 3.14, as_type="float") + self.assertAlmostEqual(frame.get_parameter("float_as_float"), 3.14, places=5) + class FrameReadTest(unittest.TestCase): """Unit tests for the Frame python bindings for Frames read from file. diff --git a/python/podio/test_utils.py b/python/podio/test_utils.py index 9a8d587ed..44efc9cce 100644 --- a/python/podio/test_utils.py +++ b/python/podio/test_utils.py @@ -43,6 +43,8 @@ def create_frame(): frame.put_parameter("an_int", 42) frame.put_parameter("some_floats", [1.23, 7.89, 3.14]) frame.put_parameter("greetings", ["from", "python"]) + frame.put_parameter("real_float", 3.14, as_type="float") + frame.put_parameter("more_real_floats", [1.23, 4.56, 7.89], as_type="float") return frame diff --git a/tests/read_python_frame.h b/tests/read_python_frame.h index 791905c18..2ec8b8090 100644 --- a/tests/read_python_frame.h +++ b/tests/read_python_frame.h @@ -76,6 +76,19 @@ int checkParameters(const podio::Frame& frame) { return 1; } + const auto realFloat = frame.getParameter("real_float"); + if (realFloat != 3.14f) { + std::cerr << "Parameter real_float was not stored correctly (expected 3.14, actual " << realFloat << ")" + << std::endl; + return 1; + } + + const auto& realFloats = frame.getParameter>("more_real_floats"); + if (realFloats.size() != 3 || realFloats[0] != 1.23 || realFloats[1] != 4.56 || realFloats[2] != 7.89) { + std::cerr << "Parameter more_real_floats was not stored as correctly (expected [1.23, 4.56, 7.89], actual" + << realFloats << ")" << std::endl; + } + return 0; }