Skip to content

Commit

Permalink
Optimize group_by when it returns objects instead of indexes #1526 (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
BoPeng authored Dec 23, 2023
1 parent b438afa commit 23bd5e9
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion src/sos/targets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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}"
Expand Down

0 comments on commit 23bd5e9

Please sign in to comment.