From 549b6a4b9f6473cb1b9bd09d21127326de3bf1cc Mon Sep 17 00:00:00 2001 From: Peter Fackeldey Date: Mon, 11 Nov 2024 15:47:13 +0100 Subject: [PATCH] fix: use None when ndims can't be inferred from a layout-like obj (#3299) --- src/awkward/_namedaxis.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/awkward/_namedaxis.py b/src/awkward/_namedaxis.py index 9f0b8d36cc..8217c30a8d 100644 --- a/src/awkward/_namedaxis.py +++ b/src/awkward/_namedaxis.py @@ -652,7 +652,7 @@ class NamedAxesWithDims: """ named_axis: list[AxisMapping] - ndims: list[int] + ndims: list[int | None] def __post_init__(self): if len(self.named_axis) != len(self.ndims): @@ -660,7 +660,7 @@ def __post_init__(self): "The number of dimensions must match the number of named axis mappings." ) - def __iter__(self) -> tp.Iterator[tuple[AxisMapping, int]]: + def __iter__(self) -> tp.Iterator[tuple[AxisMapping, int | None]]: yield from zip(self.named_axis, self.ndims) @classmethod @@ -683,18 +683,20 @@ def prepare_contexts( with HighLevelContext() as ctx: layout = ctx.unwrap(array, **_unwrap_kwargs) _named_axes.append(_get_named_axis(array)) - _ndims.append(layout.minmax_depth[1]) + _ndims.append(getattr(layout, "minmax_depth", (None, None))[1]) depth_context = {NAMED_AXIS_KEY: cls(_named_axes, _ndims)} lateral_context = {NAMED_AXIS_KEY: cls(_named_axes, _ndims)} return depth_context, lateral_context - def __setitem__(self, index: int, named_axis_with_ndim: tuple[AxisMapping, int]): + def __setitem__( + self, index: int, named_axis_with_ndim: tuple[AxisMapping, int | None] + ): named_axis, ndim = named_axis_with_ndim self.named_axis[index] = named_axis self.ndims[index] = ndim - def __getitem__(self, index: int) -> tuple[AxisMapping, int]: + def __getitem__(self, index: int) -> tuple[AxisMapping, int | None]: return self.named_axis[index], self.ndims[index] def __len__(self) -> int: