feat: Content.form_with_key_path() (#3311)

* feat: Content.form_with_key_path()

* fixed type hint (tests would have caught it)

* gitignore: add ropeproject (IDE)

* fix typing & enforce tuple record field names as strings

* style: pre-commit fixes

* forgot one type hint...

* fix typo and sort import

* add tests for  for each array type

---------

Co-authored-by: pfackeldey <fackeldey.peter@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.gitignore

@@ -50,6 +50,7 @@ GPATH
 GRTAGS
 .vscode
 pyrightconfig.json
+.ropeproject
 # ...
 
 ############################################################# LaTeX

src/awkward/contents/bitmaskedarray.py

@@ -38,7 +38,7 @@
     ToArrowOptions,
 )
 from awkward.forms.bitmaskedform import BitMaskedForm
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.index import Index
 
 if TYPE_CHECKING:
@@ -288,6 +288,16 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> BitMaskedFo
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> BitMaskedForm:
+        return self.form_cls(
+            self._mask.form,
+            self._content._form_with_key_path((*path, None)),
+            self._valid_when,
+            self._lsb_order,
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/bytemaskedarray.py

@@ -42,7 +42,7 @@
 )
 from awkward.errors import AxisError
 from awkward.forms.bytemaskedform import ByteMaskedForm
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.index import Index
 
 if TYPE_CHECKING:
@@ -218,6 +218,15 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> ByteMaskedF
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> ByteMaskedForm:
+        return self.form_cls(
+            self._mask.form,
+            self._content._form_with_key_path((*path, None)),
+            self._valid_when,
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/content.py

@@ -54,7 +54,7 @@
     TypedDict,
 )
 from awkward._util import UNSET
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.index import Index, Index64
 
 if TYPE_CHECKING:
@@ -238,6 +238,12 @@ def _form_with_key(
     ) -> Form:
         raise NotImplementedError
 
+    def form_with_key_path(self, root: FormKeyPathT = ()) -> Form:
+        return self._form_with_key_path(root)
+
+    def _form_with_key_path(self, path: FormKeyPathT) -> Form:
+        raise NotImplementedError
+
     @property
     def form_cls(self) -> type[Form]:
         raise NotImplementedError

src/awkward/contents/emptyarray.py

@@ -35,7 +35,7 @@
 )
 from awkward.errors import AxisError
 from awkward.forms.emptyform import EmptyForm
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.index import Index
 
 if TYPE_CHECKING:
@@ -118,6 +118,9 @@ def simplified(cls, *, parameters=None, backend=None):
     def _form_with_key(self, getkey: Callable[[Content], str | None]) -> EmptyForm:
         return self.form_cls(form_key=getkey(self))
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> EmptyForm:
+        return self.form_cls(form_key=repr(path))
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/indexedarray.py

@@ -39,7 +39,7 @@
     ToArrowOptions,
 )
 from awkward.errors import AxisError
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.indexedform import IndexedForm
 from awkward.index import Index
 
@@ -214,6 +214,14 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> IndexedForm
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> IndexedForm:
+        return self.form_cls(
+            self._index.form,
+            self._content._form_with_key_path((*path, None)),
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/indexedoptionarray.py

@@ -39,7 +39,7 @@
     ToArrowOptions,
 )
 from awkward.errors import AxisError
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.indexedoptionform import IndexedOptionForm
 from awkward.index import Index
 
@@ -202,6 +202,14 @@ def _form_with_key(
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> IndexedOptionForm:
+        return self.form_cls(
+            self._index.form,
+            self._content._form_with_key_path((*path, None)),
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/listarray.py

@@ -38,7 +38,7 @@
     ToArrowOptions,
 )
 from awkward.contents.listoffsetarray import ListOffsetArray
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.listform import ListForm
 from awkward.index import Index
 
@@ -207,6 +207,15 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> ListForm:
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> ListForm:
+        return self.form_cls(
+            self._starts.form,
+            self._stops.form,
+            self._content._form_with_key_path((*path, None)),
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/listoffsetarray.py

@@ -39,7 +39,7 @@
     ToArrowOptions,
 )
 from awkward.errors import AxisError
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.listoffsetform import ListOffsetForm
 from awkward.index import Index, Index64
 
@@ -199,6 +199,14 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> ListOffsetF
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> ListOffsetForm:
+        return self.form_cls(
+            self._offsets.form,
+            self._content._form_with_key_path((*path, None)),
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/numpyarray.py

@@ -45,7 +45,7 @@
     ToArrowOptions,
 )
 from awkward.errors import AxisError
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.numpyform import NumpyForm
 from awkward.index import Index
 from awkward.types.numpytype import primitive_to_dtype
@@ -200,6 +200,14 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> NumpyForm:
             form_key=getkey(self),
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> NumpyForm:
+        return self.form_cls(
+            ak.types.numpytype.dtype_to_primitive(self._data.dtype),
+            self.inner_shape,
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/recordarray.py

@@ -40,7 +40,7 @@
     ToArrowOptions,
 )
 from awkward.errors import AxisError
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.recordform import RecordForm
 from awkward.index import Index
 from awkward.record import Record
@@ -316,6 +316,22 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> RecordForm:
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> RecordForm:
+        # explicitly use `self.fields` instead of `self._fields`,
+        # because we want string-typed field names in the path -
+        # also for tuple records
+        contents = [
+            x._form_with_key_path((*path, k))
+            for k, x in zip(self.fields, self._contents)
+        ]
+
+        return self.form_cls(
+            contents,
+            self._fields,
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/regulararray.py

@@ -37,7 +37,7 @@
     RemoveStructureOptions,
     ToArrowOptions,
 )
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.regularform import RegularForm
 from awkward.index import Index
 
@@ -211,6 +211,14 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> RegularForm
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> RegularForm:
+        return self.form_cls(
+            self._content._form_with_key_path((*path, None)),
+            self._size,
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/unionarray.py

@@ -39,7 +39,7 @@
     ToArrowOptions,
 )
 from awkward.errors import AxisError
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.unionform import UnionForm
 from awkward.index import Index, Index8, Index64
 
@@ -460,6 +460,15 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> UnionForm:
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> UnionForm:
+        return self.form_cls(
+            self._tags.form,
+            self._index.form,
+            [x._form_with_key_path((*path, i)) for i, x in enumerate(self._contents)],
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/contents/unmaskedarray.py

@@ -39,7 +39,7 @@
     ToArrowOptions,
 )
 from awkward.errors import AxisError
-from awkward.forms.form import Form
+from awkward.forms.form import Form, FormKeyPathT
 from awkward.forms.unmaskedform import UnmaskedForm
 from awkward.index import Index
 
@@ -138,6 +138,13 @@ def _form_with_key(self, getkey: Callable[[Content], str | None]) -> UnmaskedFor
             form_key=form_key,
         )
 
+    def _form_with_key_path(self, path: FormKeyPathT) -> UnmaskedForm:
+        return self.form_cls(
+            self._content._form_with_key_path((*path, None)),
+            parameters=self._parameters,
+            form_key=repr(path),
+        )
+
     def _to_buffers(
         self,
         form: Form,

src/awkward/forms/form.py

@@ -27,13 +27,17 @@
     Iterator,
     JSONMapping,
     Self,
+    Tuple,
+    TypeAlias,
+    Union,
 )
 
 __all__ = ("from_dict", "from_type", "from_json", "reserved_nominal_parameters", "Form")
 
 np = NumpyMetadata.instance()
 numpy_backend = NumpyBackend.instance()
 
+FormKeyPathT: TypeAlias = Tuple[Union[str, int, None], ...]
 
 reserved_nominal_parameters: Final = frozenset(
     {