fix: more testing and label support

noxfile.py

@@ -20,7 +20,7 @@ def lint(session):
     session.run("pre-commit", "run", "--all-files", *session.posargs)
 
 
-@nox.session(python=ALL_PYTHONS)
+@nox.session(python=ALL_PYTHONS, reuse_venv=True)
 def tests(session):
     """
     Run the unit and regular tests.

pyproject.toml

@@ -5,6 +5,15 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools_scm]
 write_to = "src/hist/version.py"
 
+[tool.pytest.ini_options]
+minverison = "6.0"
+addopts = "-ra -Wd --strict-markers"
+xfail_strict = true
+testpaths = ["tests"]
+required_plugins = ["pytest-mpl"]
+log_cli_level = "DEBUG"
+
+
 [tool.nbqa.mutate]
 black = 1
 pyupgrade = 1

setup.py

@@ -18,7 +18,7 @@
 
 extras_require["test"] = [
     *extras_require["plot"],
-    "pytest >=4.6",
+    "pytest >=6",
     "pytest-mpl >=0.12",
 ]
 

src/hist/stack.py

@@ -1,12 +1,15 @@
-import sys
 import typing
-from typing import Any, Iterator, List, Tuple, Union
+from typing import Any, Iterator, List, Tuple, TypeVar, Union
 
 from .basehist import BaseHist
 
 if typing.TYPE_CHECKING:
     from mplhep.plot import Hist1DArtists
 
+__all__ = ("Stack",)
+
+T = TypeVar("T", bound="Stack")
+
 
 class Stack:
     def __init__(
@@ -30,35 +33,48 @@ def __init__(
             if first_axes != a.axes:
                 raise ValueError("The Histogram axes don't match")
 
-    def __repr__(self) -> str:
-        str_stack = ", ".join(repr(h) for h in self._stack)
-        return f"{self.__class__.__name__}({str_stack})"
+    @typing.overload
+    def __getitem__(self, val: int) -> BaseHist:
+        ...
+
+    @typing.overload
+    def __getitem__(self: T, val: slice) -> T:
+        ...
+
+    def __getitem__(self: T, val: Union[int, slice]) -> Union[BaseHist, T]:
+        if isinstance(val, slice):
+            return self.__class__(*self._stack.__getitem__(val))
 
-    def __getitem__(
-        self, val: Union[int, slice]
-    ) -> Union[BaseHist, Tuple[BaseHist, ...]]:
         return self._stack.__getitem__(val)
 
     def __iter__(self) -> Iterator[BaseHist]:
         return iter(self._stack)
 
-    def plot(self, *, overlay: None = None, **kwargs: Any) -> "List[Hist1DArtists]":
+    def __len__(self) -> int:
+        return len(self._stack)
+
+    def __repr__(self) -> str:
+        str_stack = ", ".join(repr(h) for h in self)
+        return f"{self.__class__.__name__}({str_stack})"
+
+    def plot(self, **kwargs: Any) -> "List[Hist1DArtists]":
         """
         Plot method for Stack object.
         """
-        if overlay is not None:
-            raise NotImplementedError("Currently overlay is not supported")
 
-        if self._stack[0].ndim != 1:
+        import hist.plot
+
+        if self[0].ndim != 1:
             raise NotImplementedError("Please project to 1D before calling plot")
 
-        try:
-            import mplhep.plot
-        except ModuleNotFoundError:
-            print(
-                f"{self.__class__.__name__}.plot() requires mplhep to plot, either install hist[plot] or mplhep",
-                file=sys.stderr,
-            )
-            raise
+        if "label" not in kwargs:
+            # TODO: add .name to static typing. And runtime, for that matter.
+            if all(getattr(h, "name", None) is not None for h in self):
+                kwargs["label"] = [h.name for h in self]  # type: ignore
+
+        return hist.plot.histplot(list(self), **kwargs)  # type: ignore
+
 
-        return mplhep.plot.histplot(list(self._stack), **kwargs)  # type: ignore
+# Python 3.7 only
+def __dir__() -> Tuple[str, ...]:
+    return __all__

tests/test_stacks.py

@@ -152,7 +152,7 @@ def test_stack_constructor_fails():
 
 
 @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires Python 3.7 or higher")
-def test_stack_plot():
+def test_stack_plot_construct():
     """
     Test stack plot -- whether Stack can be properly plot.
     """
@@ -217,3 +217,24 @@ def test_stack_method():
     assert s2[0].axes[0] == h.axes[0]
     assert s2[0].name == "one"
     assert s2[1].name == "two"
+
+
+def collect(*args, **kwargs):
+    return args, kwargs
+
+
+def test_stack_plot(monkeypatch):
+    import hist.plot
+
+    monkeypatch.setattr(hist.plot, "histplot", collect)
+
+    h = Hist.new.Regular(10, 0, 1).StrCategory(["one", "two"], name="str").Double()
+    s = h.stack(1)
+
+    args, kwargs = s.plot(silly=...)
+
+    assert len(s) == 2
+    assert len(list(s)) == 2
+
+    assert args == (list(s),)
+    assert kwargs == {"label": ["one", "two"], "silly": ...}