feat: add recursive argument to expected_from_buffers

src/awkward/forms/bitmaskedform.py

@@ -262,7 +262,8 @@ def __setstate__(self, state):
             )
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         yield (getkey(self, "mask"), index_to_dtype[self._mask])
-        yield from self._content._expected_from_buffers(getkey)
+        if recursive:
+            yield from self._content._expected_from_buffers(getkey, recursive)

src/awkward/forms/bytemaskedform.py

@@ -225,7 +225,8 @@ def __setstate__(self, state):
             )
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         yield (getkey(self, "mask"), index_to_dtype[self._mask])
-        yield from self._content._expected_from_buffers(getkey)
+        if recursive:
+            yield from self._content._expected_from_buffers(getkey, recursive)

src/awkward/forms/emptyform.py

@@ -151,6 +151,6 @@ def __setstate__(self, state):
             self.__init__(form_key=form_key)
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         yield from ()

src/awkward/forms/form.py

@@ -657,18 +657,25 @@ def prepare(form, multiplier):
         )
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         raise NotImplementedError
 
-    def expected_from_buffers(self, buffer_key="{form_key}-{attribute}"):
+    def expected_from_buffers(
+        self, buffer_key="{form_key}-{attribute}", recursive=True
+    ):
         """
         Args:
             buffer_key (str or callable): Python format string containing
                 `"{form_key}"` and/or `"{attribute}"` or a function that takes these
                 as keyword arguments and returns a string to use as a key for a buffer
                 in the `container`.
+            recursive (bool): If True, recurse into subforms; otherwise, yield
+                only the (buffer_key, dtype) pairs for this form object.
+
+        Yield (buffer_key, dtype) pairs describing the expected buffer keys,
+        and their corresponding dtypes, that a call to #ak.from_buffers would
+        be expected to find from the `container` object.
         """
         getkey = regularize_buffer_key(buffer_key)
-
-        return dict(self._expected_from_buffers(getkey))
+        return dict(self._expected_from_buffers(getkey, recursive))

src/awkward/forms/indexedform.py

@@ -229,7 +229,8 @@ def __setstate__(self, state):
             self.__init__(index, content, parameters=parameters, form_key=form_key)
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         yield (getkey(self, "index"), index_to_dtype[self._index])
-        yield from self._content._expected_from_buffers(getkey)
+        if recursive:
+            yield from self._content._expected_from_buffers(getkey, recursive)

src/awkward/forms/indexedoptionform.py

@@ -210,7 +210,8 @@ def __setstate__(self, state):
             self.__init__(index, content, parameters=parameters, form_key=form_key)
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         yield (getkey(self, "index"), index_to_dtype[self._index])
-        yield from self._content._expected_from_buffers(getkey)
+        if recursive:
+            yield from self._content._expected_from_buffers(getkey, recursive)

src/awkward/forms/listform.py

@@ -236,8 +236,9 @@ def __setstate__(self, state):
             )
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         yield (getkey(self, "starts"), index_to_dtype[self._starts])
         yield (getkey(self, "stops"), index_to_dtype[self._stops])
-        yield from self._content._expected_from_buffers(getkey)
+        if recursive:
+            yield from self._content._expected_from_buffers(getkey, recursive)

src/awkward/forms/listoffsetform.py

@@ -198,7 +198,8 @@ def __setstate__(self, state):
             self.__init__(offsets, content, parameters=parameters, form_key=form_key)
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         yield (getkey(self, "offsets"), index_to_dtype[self._offsets])
-        yield from self._content._expected_from_buffers(getkey)
+        if recursive:
+            yield from self._content._expected_from_buffers(getkey, recursive)

src/awkward/forms/numpyform.py

@@ -293,7 +293,7 @@ def __setstate__(self, state):
             )
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         from awkward.types.numpytype import primitive_to_dtype
 

src/awkward/forms/recordform.py

@@ -309,7 +309,8 @@ def __setstate__(self, state):
             )
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
-        for content in self._contents:
-            yield from content._expected_from_buffers(getkey)
+        if recursive:
+            for content in self._contents:
+                yield from content._expected_from_buffers(getkey, recursive)

src/awkward/forms/regularform.py

@@ -193,6 +193,7 @@ def __setstate__(self, state):
             self.__init__(content, size, parameters=parameters, form_key=form_key)
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
-        yield from self._content._expected_from_buffers(getkey)
+        if recursive:
+            yield from self._content._expected_from_buffers(getkey, recursive)

src/awkward/forms/unionform.py

@@ -295,9 +295,10 @@ def __setstate__(self, state):
             )
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
         yield (getkey(self, "tags"), index_to_dtype[self._tags])
         yield (getkey(self, "index"), index_to_dtype[self._index])
-        for content in self._contents:
-            yield from content._expected_from_buffers(getkey)
+        if recursive:
+            for content in self._contents:
+                yield from content._expected_from_buffers(getkey, recursive)

src/awkward/forms/unmaskedform.py

@@ -182,6 +182,7 @@ def __setstate__(self, state):
             self.__init__(content, parameters=parameters, form_key=form_key)
 
     def _expected_from_buffers(
-        self, getkey: Callable[[Form, str], str]
+        self, getkey: Callable[[Form, str], str], recursive: bool
     ) -> Iterator[tuple[str, np.dtype]]:
-        yield from self._content._expected_from_buffers(getkey)
+        if recursive:
+            yield from self._content._expected_from_buffers(getkey, recursive)

tests/test_2724_expected_from_buffers_recursive.py

@@ -0,0 +1,29 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+import pytest  # noqa: F401
+
+import awkward as ak
+
+# 6.6 and 7.7 are inaccessible
+layout = ak.contents.listoffsetarray.ListOffsetArray(
+    ak.index.Index(np.array([1, 4, 4, 6], dtype=np.int64)),
+    ak.contents.numpyarray.NumpyArray(
+        np.array([6.6, 1.1, 2.2, 3.3, 4.4, 5.5, 7.7], dtype=np.float64)
+    ),
+)
+
+
+def test_recursive():
+    form, length, container = ak.to_buffers(layout)
+    assert form.expected_from_buffers(recursive=True) == {
+        "node0-offsets": np.dtype("int64"),
+        "node1-data": np.dtype("float64"),
+    }
+
+
+def test_non_recursive():
+    form, length, container = ak.to_buffers(layout)
+    assert form.expected_from_buffers(recursive=False) == {
+        "node0-offsets": np.dtype("int64")
+    }