From 23bd5e95f977260036b23e19a7cd782548bc8ce3 Mon Sep 17 00:00:00 2001 From: Bo Date: Sat, 23 Dec 2023 00:19:23 -0600 Subject: [PATCH] Optimize group_by when it returns objects instead of indexes #1526 (#1529) --- src/sos/targets.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/sos/targets.py b/src/sos/targets.py index 7fe653466..afb253e69 100644 --- a/src/sos/targets.py +++ b/src/sos/targets.py @@ -1772,6 +1772,7 @@ def _group(self, by): raise ValueError( f"Customized grouping method should return a list. {idx} of type {idx.__class__.__name__} is returned." ) from e + _target_index = None for grp in by(self): if isinstance(grp, Sequence) and all( isinstance(x, int) for x in grp): @@ -1784,7 +1785,9 @@ def _group(self, by): index = [] for x in sos_targets(grp): try: - index.append(self._targets.index(x)) + if _target_index is None: + _target_index = {v:k for k,v in enumerate(self._targets)} + index.append(_target_index[x]) except Exception as e: raise ValueError( f"Returned target is not one of the targets. {x}"