Merge pull request #448 from dask-contrib/agoose77/fix-attrs-readonly

fix: return proxy of attrs, behavior; dak.from_awkward correctly sets attrs

src/dask_awkward/lib/core.py

@@ -11,6 +11,7 @@
 from functools import cached_property, partial, wraps
 from inspect import getattr_static
 from numbers import Number
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Literal, TypeVar, Union, overload
 
 import awkward as ak
@@ -1090,17 +1091,18 @@ def layout(self) -> Content:
         raise ValueError("This collection's meta is None; unknown layout.")
 
     @property
-    def attrs(self) -> dict:
+    def attrs(self) -> Mapping[str, Any]:
         """awkward Array attrs dictionary."""
         if self._meta is not None:
-            return self._meta.attrs
+            return MappingProxyType(self._meta.attrs)
         raise ValueError("This collection's meta is None; no attrs property available.")
 
     @property
-    def behavior(self) -> Mapping:
+    def behavior(self) -> Mapping | None:
         """awkward Array behavior dictionary."""
         if self._meta is not None:
-            return self._meta.behavior
+            behavior = self._meta.behavior
+            return None if behavior is None else MappingProxyType(behavior)
         raise ValueError(
             "This collection's meta is None; no behavior property available."
         )

src/dask_awkward/lib/io/columnar.py

@@ -41,6 +41,10 @@ def form(self) -> Form:
     def behavior(self) -> dict | None:
         ...
 
+    @property
+    def attrs(self) -> dict | None:
+        ...
+
     def project_columns(self: T, columns: frozenset[str]) -> T:
         ...
 
@@ -62,7 +66,8 @@ class ColumnProjectionMixin(ImplementsNecessaryColumns[FormStructure]):
 
     def mock(self: S) -> AwkwardArray:
         return cast(
-            AwkwardArray, typetracer_from_form(self.form, behavior=self.behavior)
+            AwkwardArray,
+            typetracer_from_form(self.form, behavior=self.behavior, attrs=self.attrs),
         )
 
     def mock_empty(self: S, backend: BackendT = "cpu") -> AwkwardArray:

src/dask_awkward/lib/io/io.py

@@ -55,8 +55,8 @@ def __init__(
     ) -> None:
         self.arr = arr
         self.form = arr.layout.form
-        self.behavior = behavior
-        self.attrs = attrs
+        self.behavior = behavior if behavior else arr.behavior
+        self.attrs = attrs if attrs else arr.attrs
 
     @property
     def use_optimization(self):
@@ -119,7 +119,7 @@ def from_awkward(
     return cast(
         Array,
         from_map(
-            FromAwkwardFn(source, behavior=behavior),
+            FromAwkwardFn(source, behavior=behavior, attrs=attrs),
             starts_stops,
             label=label or "from-awkward",
             token=tokenize(source, npartitions),

src/dask_awkward/lib/io/json.py

@@ -49,6 +49,7 @@ def __init__(
         compression: str | None = None,
         schema: str | dict | list | None = None,
         behavior: Mapping | None = None,
+        attrs: Mapping[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         self.compression = compression
@@ -57,6 +58,7 @@ def __init__(
         self.kwargs = kwargs
         self.form = form
         self.behavior = behavior
+        self.attrs = attrs
 
     @abc.abstractmethod
     def __call__(self, source: Any) -> ak.Array:
@@ -97,6 +99,7 @@ def __init__(
         compression: str | None = None,
         schema: str | dict | list | None = None,
         behavior: Mapping | None = None,
+        attrs: Mapping[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -105,6 +108,7 @@ def __init__(
             schema=schema,
             form=form,
             behavior=behavior,
+            attrs=attrs,
             **kwargs,
         )
 
@@ -131,6 +135,7 @@ def __init__(
         compression: str | None = None,
         schema: str | dict | list | None = None,
         behavior: Mapping | None = None,
+        attrs: Mapping[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -139,6 +144,7 @@ def __init__(
             schema=schema,
             form=form,
             behavior=behavior,
+            attrs=attrs,
             **kwargs,
         )
 
@@ -169,13 +175,15 @@ def __init__(
         compression: str | None = None,
         schema: str | dict | list | None = None,
         behavior: Mapping | None = None,
+        attrs: Mapping[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
             storage=storage,
             compression=compression,
             schema=schema,
             behavior=behavior,
+            attrs=attrs,
             form=form,
             **kwargs,
         )

tests/test_core.py

@@ -891,6 +891,28 @@ def test_shape_only_ops(fn: Callable, tmp_path_factory: pytest.TempPathFactory)
         result.compute()
 
 
+def test_assign_behavior() -> None:
+    behavior = {"test": "hello"}
+    x = ak.Array([{"a": 1, "b": 2}, {"a": 3, "b": 4}], behavior=behavior, attrs={})
+    dx = dak.from_awkward(x, 3)
+    with pytest.raises(
+        TypeError, match="'mappingproxy' object does not support item assignment"
+    ):
+        dx.behavior["should_fail"] = None  # type: ignore
+    assert dx.behavior == behavior
+
+
+def test_assign_attrs() -> None:
+    attrs = {"test": "hello"}
+    x = ak.Array([{"a": 1, "b": 2}, {"a": 3, "b": 4}], behavior={}, attrs=attrs)
+    dx = dak.from_awkward(x, 3)
+    with pytest.raises(
+        TypeError, match="'mappingproxy' object does not support item assignment"
+    ):
+        dx.attrs["should_fail"] = None  # type: ignore
+    assert dx.attrs == attrs
+
+
 @delayed
 def a_delayed_array():
     return ak.Array([2, 4])