fix: support pick_set inside a larger expression (#793)

* tests: adding pick inside a larger expression

Signed-off-by: Henry Schreiner <[email protected]>

* fix: allow picking and pick_set to coexist

Signed-off-by: Henry Schreiner <[email protected]>

Signed-off-by: Henry Schreiner <[email protected]>

Author: henryiii

src/boost_histogram/_internal/hist.py

@@ -927,6 +927,9 @@ def __getitem__(  # noqa: C901
             new_reduced.view(flow=True)[...] = reduced.view(flow=True)[tuple_slice]
             reduced = new_reduced
             integrations = {i - sum(j <= i for j in pick_each) for i in integrations}
+            pick_set = {
+                i - sum(j <= i for j in pick_each): v for i, v in pick_set.items()
+            }
             for slice_ in slices:
                 slice_.iaxis -= sum(j <= slice_.iaxis for j in pick_each)
 
@@ -945,9 +948,8 @@ def __getitem__(  # noqa: C901
                 selection = copy.copy(pick_set[i])
                 ax = reduced.axis(i)
                 if ax.traits_ordered:
-                    raise RuntimeError(
-                        f"Axis {i} is not a categorical axis, cannot pick with list"
-                    )
+                    msg = f"Axis {i} is not a categorical axis, cannot pick with list: {ax}"
+                    raise RuntimeError(msg)
 
                 if ax.traits_overflow and ax.size not in pick_set[i]:
                     selection.append(ax.size)

tests/test_internal_histogram.py

@@ -107,6 +107,27 @@ def test_str_categories_histogram():
     assert hist[bh.loc("c")] == 1
 
 
+# Issue 715
+
+
+def test_select_many():
+    hist = bh.Histogram(
+        bh.axis.StrCategory(["a", "b"]),
+        bh.axis.StrCategory(["x", "y", "z"]),
+        bh.axis.Regular(10, 0, 1),
+    )
+
+    pick_a = hist[bh.loc("a"), ...]
+    with pytest.warns(UserWarning):
+        pick_b = pick_a[[bh.loc("x"), bh.loc("y")], ...]
+
+    with pytest.warns(UserWarning):
+        pick = hist[bh.loc("a"), [bh.loc("x"), bh.loc("y")], ...]
+
+    assert pick_b.axes[0] == pick.axes[0]
+    assert len(pick_b.axes) == len(pick.axes)
+
+
 def test_growing_histogram():
     hist = bh.Histogram(
         bh.axis.Regular(10, 0, 1, growth=True), storage=bh.storage.Int64()