feat: axis sort

src/hist/basehist.py

@@ -28,7 +28,7 @@
 from .quick_construct import MetaConstructor
 from .storage import Storage
 from .svgplots import html_hist, svg_hist_1d, svg_hist_1d_c, svg_hist_2d, svg_hist_nd
-from .typing import ArrayLike, SupportsIndex
+from .typing import ArrayLike, Protocol, SupportsIndex
 
 if typing.TYPE_CHECKING:
     from builtins import ellipsis
@@ -38,6 +38,12 @@
 
     from .plot import FitResultArtists, MainAxisArtists, RatiolikeArtists
 
+
+class SupportsLessThan(Protocol):
+    def __lt__(self, __other: Any) -> bool:
+        ...
+
+
 InnerIndexing = Union[
     SupportsIndex, str, Callable[[bh.axis.Axis], int], slice, "ellipsis"
 ]
@@ -228,6 +234,24 @@ def fill(
         total_data = tuple(args) + tuple(data)  # Python 2 can't unpack twice
         return super().fill(*total_data, weight=weight, sample=sample, threads=threads)
 
+    def sort(
+        self: T,
+        axis: Union[int, str],
+        key: Union[
+            Callable[[int], SupportsLessThan], Callable[[str], SupportsLessThan], None
+        ] = None,
+        reverse: bool = False,
+    ) -> T:
+        """
+        Sort a categorical axis.
+        """
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            sorted_cats = sorted(self.axes[axis], key=key, reverse=reverse)
+            # This can only return T, not float, etc., so we ignore the extra types here
+            return self[{axis: [bh.loc(x) for x in sorted_cats]}]  # type: ignore
+
     def _loc_shortcut(self, x: Any) -> Any:
         """
         Convert some specific indices to location.

tests/test_general.py

@@ -896,3 +896,11 @@ def test_select_by_index_imag():
 
     assert tuple(h[[2, 1]].axes[0]) == (9, 8)
     assert tuple(h[[8j, 7j]].axes[0]) == (8, 7)
+
+
+def test_sorted_simple():
+    h = Hist.new.IntCat([4, 1, 2]).Double()
+    assert tuple(h.sort(0).axes[0]) == (1, 2, 4)
+    assert tuple(h.sort(0, reverse=True).axes[0]) == (4, 2, 1)
+    assert tuple(h.sort(0, key=lambda x: -x).axes[0]) == (4, 2, 1)
+    assert tuple(h.sort(0, key=lambda x: -x, reverse=True).axes[0]) == (1, 2, 4)