From 0a75ed166dceb7958a9eacd2b20159238b8936e8 Mon Sep 17 00:00:00 2001 From: Bo Peng Date: Mon, 18 Dec 2023 01:17:39 -0600 Subject: [PATCH] Optimize group_by when it returns objects instead of indexes #1526 --- 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}"