Merge branch 'main' into jpivarski/fix-json-from-schema-with-missing-…

…field

.github/workflows/deploy-cpp.yml

@@ -91,7 +91,7 @@ jobs:
     - name: Prepare build files
       run: pipx run nox -s prepare
 
-    - uses: pypa/cibuildwheel@v2.15.0
+    - uses: pypa/cibuildwheel@v2.16.0
       env:
         CIBW_BUILD: ${{ matrix.build }}
         CIBW_ARCHS: ${{ matrix.arch }}
@@ -140,7 +140,7 @@ jobs:
 
     - uses: docker/setup-qemu-action@v3.0.0
 
-    - uses: pypa/cibuildwheel@v2.15.0
+    - uses: pypa/cibuildwheel@v2.16.0
       env:
         CIBW_BUILD: cp${{ matrix.python }}-*
         CIBW_ARCHS: ${{ matrix.arch }}

.github/workflows/docs.yml

@@ -319,7 +319,7 @@ jobs:
       uses: aws-actions/configure-aws-credentials@v4
       with:
         aws-region: eu-west-2
-        role-to-assume: ${{ secrets.AWS_DEPLOY_ROLE }}
+        role-to-assume: arn:aws:iam::${{ secrets.AWS_ACCOUNT_ID }}:role/${{ secrets.AWS_DEPLOY_ROLE }}
     - name: Download rendered docs
       uses: actions/download-artifact@v3
       with:
@@ -351,7 +351,7 @@ jobs:
         uses: aws-actions/configure-aws-credentials@v4
         with:
           aws-region: eu-west-2
-          role-to-assume: ${{ secrets.AWS_DEPLOY_ROLE }}
+          role-to-assume: arn:aws:iam::${{ secrets.AWS_ACCOUNT_ID }}:role/${{ secrets.AWS_DEPLOY_ROLE }}
       - name: Download rendered docs
         uses: actions/download-artifact@v3
         with:

.github/workflows/packaging-test.yml

@@ -56,15 +56,15 @@ jobs:
     - name: Prepare build files
       run: pipx run nox -s prepare
 
-    - uses: pypa/cibuildwheel@v2.15.0
+    - uses: pypa/cibuildwheel@v2.16.0
       env:
         CIBW_ARCHS_MACOS: universal2
         CIBW_BUILD: cp39-win_amd64 cp310-manylinux_x86_64 cp38-macosx_universal2
       with:
         config-file: cibuildwheel.toml
         package-dir: awkward-cpp
 
-    - uses: pypa/cibuildwheel@v2.15.0
+    - uses: pypa/cibuildwheel@v2.16.0
       if: matrix.os == 'ubuntu-latest'
       env:
         CIBW_BUILD: cp312-manylinux_x86_64

.github/workflows/semantic-pr-title.yml

@@ -17,6 +17,6 @@ jobs:
     name: Validate PR title
     runs-on: ubuntu-latest
     steps:
-      - uses: amannn/action-semantic-pull-request@v5.2.0
+      - uses: amannn/action-semantic-pull-request@v5.3.0
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.pre-commit-config.yaml

@@ -39,7 +39,7 @@ repos:
     additional_dependencies: [pyyaml]
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.0.290
+  rev: v0.0.291
   hooks:
   - id: ruff
     args: ["--fix", "--show-fixes"]
@@ -74,6 +74,6 @@ repos:
 
 
 - repo: https://github.com/asottile/pyupgrade
-  rev: v3.11.0
+  rev: v3.13.0
   hooks:
   - id: pyupgrade

src/awkward/_nplikes/typetracer.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 from numbers import Number
+from typing import Callable
 
 import numpy
 
@@ -610,7 +611,7 @@ def populate_shape_and_items(node, dim):
                     if has_seen_leaf:
                         if len(node) != shape[dim - 1]:
                             raise ValueError(
-                                f"sequence at dimension {dim} does not match shape {shape[dim-1]}"
+                                f"sequence at dimension {dim} does not match shape {shape[dim - 1]}"
                             )
                     else:
                         shape.append(len(node))
@@ -1394,70 +1395,75 @@ def __dlpack__(self, stream=None):
 
 
 def _attach_report(
-    layout: ak.contents.Content, form: ak.forms.Form, report: TypeTracerReport
+    layout: ak.contents.Content,
+    form: ak.forms.Form,
+    report: TypeTracerReport,
+    getkey: Callable[[ak.forms.form, str], str],
 ):
     if isinstance(layout, (ak.contents.BitMaskedArray, ak.contents.ByteMaskedArray)):
         assert isinstance(form, (ak.forms.BitMaskedForm, ak.forms.ByteMaskedForm))
-        layout.mask.data.form_key = form.form_key
+        layout.mask.data.form_key = getkey(form, "mask")
         layout.mask.data.report = report
-        _attach_report(layout.content, form.content, report)
+        _attach_report(layout.content, form.content, report, getkey)
 
     elif isinstance(layout, ak.contents.EmptyArray):
         assert isinstance(form, ak.forms.EmptyForm)
 
     elif isinstance(layout, (ak.contents.IndexedArray, ak.contents.IndexedOptionArray)):
         assert isinstance(form, (ak.forms.IndexedForm, ak.forms.IndexedOptionForm))
-        layout.index.data.form_key = form.form_key
+        layout.index.data.form_key = getkey(form, "index")
         layout.index.data.report = report
-        _attach_report(layout.content, form.content, report)
+        _attach_report(layout.content, form.content, report, getkey)
 
     elif isinstance(layout, ak.contents.ListArray):
         assert isinstance(form, ak.forms.ListForm)
-        layout.starts.data.form_key = form.form_key
+        layout.starts.data.form_key = getkey(form, "starts")
         layout.starts.data.report = report
-        layout.stops.data.form_key = form.form_key
+        layout.stops.data.form_key = getkey(form, "stops")
         layout.stops.data.report = report
-        _attach_report(layout.content, form.content, report)
+        _attach_report(layout.content, form.content, report, getkey)
 
     elif isinstance(layout, ak.contents.ListOffsetArray):
         assert isinstance(form, ak.forms.ListOffsetForm)
-        layout.offsets.data.form_key = form.form_key
+        layout.offsets.data.form_key = getkey(form, "offsets")
         layout.offsets.data.report = report
-        _attach_report(layout.content, form.content, report)
+        _attach_report(layout.content, form.content, report, getkey)
 
     elif isinstance(layout, ak.contents.NumpyArray):
         assert isinstance(form, ak.forms.NumpyForm)
-        layout.data.form_key = form.form_key
+        layout.data.form_key = getkey(form, "data")
         layout.data.report = report
 
     elif isinstance(layout, ak.contents.RecordArray):
         assert isinstance(form, ak.forms.RecordForm)
         for x, y in zip(layout.contents, form.contents):
-            _attach_report(x, y, report)
+            _attach_report(x, y, report, getkey)
 
     elif isinstance(layout, (ak.contents.RegularArray, ak.contents.UnmaskedArray)):
         assert isinstance(form, (ak.forms.RegularForm, ak.forms.UnmaskedForm))
-        _attach_report(layout.content, form.content, report)
+        _attach_report(layout.content, form.content, report, getkey)
 
     elif isinstance(layout, ak.contents.UnionArray):
         assert isinstance(form, ak.forms.UnionForm)
-        layout.tags.data.form_key = form.form_key
+        layout.tags.data.form_key = getkey(form, "tags")
         layout.tags.data.report = report
-        layout.index.data.form_key = form.form_key
+        layout.index.data.form_key = getkey(form, "index")
         layout.index.data.report = report
         for x, y in zip(layout.contents, form.contents):
-            _attach_report(x, y, report)
+            _attach_report(x, y, report, getkey)
 
     else:
         raise AssertionError(f"unrecognized layout type {type(layout)}")
 
 
 def typetracer_with_report(
-    form: ak.forms.Form, forget_length: bool = True
+    form: ak.forms.Form,
+    getkey: Callable[[ak.forms.form, str], str],
+    forget_length: bool = True,
 ) -> tuple[ak.contents.Content, TypeTracerReport]:
     layout = form.length_zero_array(highlevel=False).to_typetracer(
         forget_length=forget_length
     )
     report = TypeTracerReport()
-    _attach_report(layout, form, report)
+    _attach_report(layout, form, report, getkey)
     return layout, report

src/awkward/typetracer.py

@@ -30,6 +30,7 @@
 from awkward._util import UNSET
 from awkward.contents import Content
 from awkward.forms import Form
+from awkward.forms.form import regularize_buffer_key
 from awkward.highlevel import Array, Record
 from awkward.operations.ak_to_layout import to_layout
 from awkward.types.numpytype import is_primitive
@@ -100,12 +101,21 @@ def touch_data(array, *, highlevel: bool = True, behavior=None) -> T:
 
 
 def typetracer_with_report(
-    form, forget_length: bool = UNSET, *, highlevel: bool = False, behavior=None
+    form,
+    forget_length: bool = UNSET,
+    *,
+    buffer_key="{form_key}",
+    highlevel: bool = False,
+    behavior=None,
 ) -> tuple[Content, TypeTracerReport]:
     """
     Args:
         form (#ak.forms.Form or str/dict equivalent): The form of the Awkward
             Array to build a typetracer-backed array from.
+        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 `form_key`
+            for low-level typetracer buffers.
         highlevel (bool): If True, return an #ak.Array; otherwise, return
             a low-level #ak.contents.Content subclass.
         behavior (None or dict): Custom #ak.behavior for the output array, if
@@ -135,7 +145,12 @@ def typetracer_with_report(
             "always forget lengths",
             "2.5.0",
         )
-    layout, report = _typetracer_with_report(form, forget_length=forget_length)
+
+    getkey = regularize_buffer_key(buffer_key)
+
+    layout, report = _typetracer_with_report(
+        form, getkey=getkey, forget_length=forget_length
+    )
     return wrap_layout(layout, behavior=behavior, highlevel=highlevel), report
 
 

tests/test_2719_typetracer_buffer_key.py

@@ -0,0 +1,53 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+
+import pytest  # noqa: F401
+
+import awkward as ak
+
+form_dict = {
+    "class": "RecordArray",
+    "fields": ["x"],
+    "contents": [
+        {
+            "class": "ListArray",
+            "starts": "i64",
+            "stops": "i64",
+            "content": {
+                "class": "NumpyArray",
+                "primitive": "int64",
+                "inner_shape": [],
+                "parameters": {},
+                "form_key": "x.list.content",
+            },
+            "parameters": {},
+            "form_key": "x.list",
+        }
+    ],
+    "parameters": {},
+}
+
+
+def test_without_attribute():
+    form = ak.forms.from_dict(form_dict)
+    array, report = ak.typetracer.typetracer_with_report(
+        form, buffer_key="{form_key}", highlevel=True
+    )
+    ak.typetracer.touch_data(array)
+    assert set(report.data_touched) == {
+        "x.list",
+        "x.list.content",
+    }
+
+
+def test_with_attribute():
+    form = ak.forms.from_dict(form_dict)
+    array, report = ak.typetracer.typetracer_with_report(
+        form, buffer_key="{form_key}-{attribute}", highlevel=True
+    )
+    ak.typetracer.touch_data(array)
+    assert set(report.data_touched) == {
+        "x.list-starts",
+        "x.list-stops",
+        "x.list.content-data",
+    }