feat: allow shortcut rebin and sum

Author: henryiii

src/boost_histogram/_internal/hist.py

@@ -442,14 +442,13 @@ def _compute_commonindex(self, index):
 
         # Allow [bh.loc(...)] to work
         for i in range(len(indexes)):
-            if callable(indexes[i]):
+            # Support sum and rebin directly
+            if indexes[i] is sum or hasattr(indexes[i], "factor"):
+                indexes[i] = slice(None, None, indexes[i])
+            # General locators
+            elif callable(indexes[i]):
                 indexes[i] = indexes[i](self.axes[i])
-            elif hasattr(indexes[i], "flow"):
-                if indexes[i].flow == 1:
-                    indexes[i] = hist.axis(i).size
-                elif indexes[i].flow == -1:
-                    indexes[i] = -1
-            elif isinstance(indexes[i], int):
+            elif hasattr(indexes[i], "__index__"):
                 if abs(indexes[i]) >= hist.axis(i).size:
                     raise IndexError("histogram index is out of range")
                 indexes[i] %= hist.axis(i).size

tests/test_histogram.py

@@ -569,9 +569,13 @@ def test_shrink_1d():
 def test_rebin_1d():
     h = bh.Histogram(bh.axis.Regular(20, 1, 5))
     h.fill(1.1)
+
     hs = h[{0: slice(None, None, bh.rebin(4))}]
     assert_array_equal(hs.view(), [1, 0, 0, 0, 0])
 
+    hs = h[{0: bh.rebin(4)}]
+    assert_array_equal(hs.view(), [1, 0, 0, 0, 0])
+
 
 def test_shrink_rebin_1d():
     h = bh.Histogram(bh.axis.Regular(20, 0, 4))
@@ -594,6 +598,7 @@ def test_rebin_nd():
     assert h[{0: s[:: bh.rebin(2)], 2: s[:: bh.rebin(2)]}].axes.size == (10, 30, 20)
 
     assert h[{1: s[:: bh.sum]}].axes.size == (20, 40)
+    assert h[{1: bh.sum}].axes.size == (20, 40)
 
 
 # CLASSIC: This used to have metadata too, but that does not compare equal

tests/test_histogram_indexing.py

@@ -73,6 +73,10 @@ def test_get_1D_slice():
 
     assert len(h1[2:4].view()) == 2
     assert len(h1[2 : 4 : bh.rebin(2)].view()) == 1
+    assert len(h1[:: bh.rebin(2)].view()) == 5
+
+    # Shortcut
+    assert len(h1[bh.rebin(2)].view()) == 5
 
     assert h1[2:4].metadata == {"that": 3}
 
@@ -123,6 +127,11 @@ def test_basic_projection():
     assert h1 == h2[..., ::sum, ::sum]
     assert h2.sum(flow=True) == h2[::sum, ::sum, ::sum]
 
+    # Shortcut
+    assert h1 == h2[:, sum, sum]
+    assert h1 == h2[..., sum, sum]
+    assert h2.sum(flow=True) == h2[sum, sum, sum]
+
 
 def test_slicing_projection():
     h1 = bh.Histogram(
@@ -192,6 +201,9 @@ def test_mix_value_with_slice_2():
     assert_array_equal(h[:, :, True].view(), vals)
     assert_array_equal(h[:, :, False].view(), 0)
 
+    h2 = h[bh.rebin(2), bh.rebin(5), :]
+    assert_array_equal(h2.shape, (5, 2, 2))
+
 
 def test_repr():
     assert repr(bh.loc(2)) == "loc(2)"