feat: add support for std::bitset (#1182)

* feat: add support for bitset

* style: pre-commit fixes

* Update test file name to match pr nr

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jim Pivarski <[email protected]>

Author: 3 people

src/uproot/containers.py

@@ -864,7 +864,9 @@ class AsVectorLike(AsContainer):
 
     def __init__(self, header, values):
         self.header = header
-        if isinstance(values, AsContainer):
+        if isinstance(self, uproot.containers.AsBitSet):
+            self._items = numpy.dtype(numpy.bool_)
+        elif isinstance(values, AsContainer):
             self._items = values
         elif isinstance(values, type) and issubclass(
             values, (uproot.model.Model, uproot.model.DispatchByVersion)
@@ -1117,7 +1119,6 @@ class AsList(AsVectorLike):
         header (bool): Sets the :ref:`uproot.containers.AsContainer.header`.
         values (:doc:`uproot.model.Model` or :doc:`uproot.containers.Container`): Data
             type for data nested in the container.
-
     A :doc:`uproot.containers.AsContainer` for ``std::list``.
     """
 
@@ -1139,6 +1140,36 @@ def _container_type(self):
         return STLList
 
 
+class AsBitSet(AsVectorLike):
+    """
+    Args:
+        header (bool): Sets the :ref:`uproot.containers.AsContainer.header`.
+        keys (:doc:`uproot.model.Model` or :doc:`uproot.containers.Container`): Data
+            type for data nested in the container.
+
+    A :doc:`uproot.containers.AsContainer` for ``std::bitset``.
+    """
+
+    _specialpathitem_name = "bitset"
+
+    @property
+    def typename(self):
+        return f"std::bitset<{_content_typename(self.keys)}>"
+
+    @property
+    def keys(self):
+        """
+        Data type for data nested in the container.
+        """
+        return self._items
+
+    _form_parameters = {"__array__": "bitset"}
+
+    @property
+    def _container_type(self):
+        return STLBitSet
+
+
 class AsSet(AsVectorLike):
     """
     Args:
@@ -1633,6 +1664,49 @@ def tolist(self):
         ]
 
 
+class STLBitSet(Container, Sequence):
+    """
+    Args:
+        values (``numpy.ndarray`` or iterable): Contents of the ``std::vector``.
+
+    Representation of a C++ ``std::bitset`` as a Python ``Sequence``.
+    """
+
+    def __init__(self, numbytes):
+        self._numbytes = numpy.asarray(numbytes)
+
+    def __str__(self, limit=85):
+        def tostring(i):
+            return _tostring(self._values[i])
+
+        return _str_with_ellipsis(tostring, len(self), "[", "]", limit)
+
+    def __repr__(self, limit=85):
+        return f"<STLBitSet {self.__str__(limit=limit - 30)} at 0x{id(self):012x}>"
+
+    def __getitem__(self, where):
+        return self._numbytes[where]
+
+    def __len__(self):
+        return self._numbytes
+
+    def __iter__(self):
+        return iter(self._numbytes)
+
+    def __eq__(self, other):
+        if isinstance(other, STLBitSet):
+            return self._numbytes == other._numbytes
+        elif isinstance(other, Sequence):
+            return self._numbytes == other
+        else:
+            return False
+
+    def tolist(self):
+        return [
+            x.tolist() if isinstance(x, (Container, numpy.ndarray)) else x for x in self
+        ]
+
+
 class STLSet(Container, Set):
     """
     Args:

src/uproot/interpretation/identify.py

@@ -964,7 +964,7 @@ def _parse_node(tokens, i, typename, file, quote, header, inner_header):
         return (
             i + 4,
             _parse_maybe_quote(
-                f'uproot.containers.AsFIXME("std::bitset<{num_bits}>")',
+                f"uproot.containers.AsBitSet({header}, {num_bits})",
                 quote,
             ),
         )

tests/test_1182_add_support_for_bitset.py

@@ -0,0 +1,40 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE
+
+import uproot
+import skhep_testdata
+
+
+def test_add_support_for_bitset():
+    with uproot.open(skhep_testdata.data_path("uproot-issue-40.root")) as f:
+        assert (
+            repr(f["tree"]["bitset8"].interpretation)
+            == "AsObjects(AsBitSet(True, dtype('bool')))"
+        )
+        assert (
+            repr(f["tree"]["bitset16"].interpretation)
+            == "AsObjects(AsBitSet(True, dtype('bool')))"
+        )
+
+        assert f["tree"]["bitset8"].array().tolist() == [
+            [True, False, True, False, True, False, False, False]
+        ]
+        assert f["tree"]["bitset16"].array().tolist() == [
+            [
+                False,
+                True,
+                False,
+                True,
+                False,
+                True,
+                True,
+                True,
+                True,
+                False,
+                True,
+                False,
+                True,
+                False,
+                False,
+                False,
+            ]
+        ]