From 29dc47f6d30b78b6460c15b0cf65f63ef5856e8d Mon Sep 17 00:00:00 2001 From: "askerosted@gmail.com" Date: Wed, 21 Feb 2024 13:08:19 +0900 Subject: [PATCH] extractor_name option to generic extractor + bug fix --- .../data/extractors/icecube/i3genericextractor.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/graphnet/data/extractors/icecube/i3genericextractor.py b/src/graphnet/data/extractors/icecube/i3genericextractor.py index e907181d0..c79b7329b 100644 --- a/src/graphnet/data/extractors/icecube/i3genericextractor.py +++ b/src/graphnet/data/extractors/icecube/i3genericextractor.py @@ -45,6 +45,7 @@ def __init__( self, keys: Optional[Union[str, List[str]]] = None, exclude_keys: Optional[Union[str, List[str]]] = None, + extractor_name: str = GENERIC_EXTRACTOR_NAME, ): """Construct I3GenericExtractor. @@ -72,7 +73,7 @@ def __init__( self._exclude_keys: Optional[List[str]] = exclude_keys # Base class constructor - super().__init__(GENERIC_EXTRACTOR_NAME) + super().__init__(extractor_name) def _get_keys(self, frame: "icetray.I3Frame") -> List[str]: """Get the list of keys to be queried from `frame`. @@ -170,6 +171,12 @@ def __call__(self, frame: "icetray.I3Frame") -> Dict[str, Any]: # Flatten all other objects else: results[key] = self._flatten_result(result) + if ( + isinstance(results[key], dict) + and "value" in results[key] + and len(results[key]) == 1 + ): + results[key] = results[key]["value"] # Serialise list of iterables to JSON results = {key: serialise(value) for key, value in results.items()}