diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index 02877d6b..9e7bb799 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -980,6 +980,8 @@ def _call_behavior_method(self, method_name: str, *args: Any, **kwargs: Any) -> themethod = getattr(self._meta, method_name) thesig = inspect.signature(themethod) if "_dask_array_" in thesig.parameters: + if "_dask_array_" not in kwargs: + kwargs["_dask_array_"] = self return themethod(*args, **kwargs) return self.map_partitions( _BehaviorMethodFn(method_name, **kwargs), diff --git a/tests/test_behavior.py b/tests/test_behavior.py index 25e76856..ea090335 100644 --- a/tests/test_behavior.py +++ b/tests/test_behavior.py @@ -23,6 +23,13 @@ def x2(self): def point_abs(self): return np.sqrt(self.x**2 + self.y**2) + @property + def non_dask_property(self, _dask_array_=None): + return "this is a non-dask property" + + def non_dask_method(self, _dask_array_=None): + return _dask_array_ + def test_distance_behavior( daa_p1: dak.Array, @@ -45,6 +52,10 @@ def test_property_behavior(daa_p1: dak.Array, caa_p1: ak.Array) -> None: assert daa.behavior == caa.behavior + assert daa.non_dask_property == caa.non_dask_property + + assert repr(daa.non_dask_method()) == repr(daa) + def test_nonexistent_behavior(daa_p1: dak.Array, daa_p2: dak.Array) -> None: daa1 = dak.with_name(daa_p1["points"], "Point", behavior=behaviors)