Merge pull request #228 from lgray/behavior_tweak

Substitute in self for _dask_array_ when it is not specified

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),

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)