typing

src/dask_histogram/boost.py

@@ -111,7 +111,7 @@ def concrete_fill(
             )
         return super().fill(*args, weight=weight, sample=sample, threads=threads)
 
-    def fill(
+    def fill(  # type: ignore
         self,
         *args: DaskCollection,
         weight: DaskCollection | None = None,
@@ -193,7 +193,7 @@ def fill(
         if self._staged is not None:
             self._staged += new_fill
         else:
-            self._staged = new_fill
+            self._staged = new_fill  # type: ignore
 
         return self
 
@@ -360,8 +360,8 @@ def to_dask_array(
             counts = da.from_array(counts)
             edges = [da.from_array(ea) for ea in edges]  # type: ignore
             if dd:
-                return counts, edges
-            return tuple([counts, *edges])
+                return counts, edges  # type: ignore
+            return tuple([counts, *edges])  # type: ignore
 
 
 def histogramdd(
@@ -375,7 +375,7 @@ def histogramdd(
     histogram: Any | None = None,
     storage: storage.Storage = storage.Double(),
     threads: int | None = None,
-) -> (Histogram | tuple[da.Array, ...] | tuple[da.Array, tuple[da.Array, ...]]):
+) -> Histogram | tuple[da.Array, ...] | tuple[da.Array, list[da.Array]]:
     """Histogram Dask data in multiple dimensions.
 
     Parameters
@@ -697,7 +697,7 @@ def histogram2d(
 
     if histogram != Histogram:
         return hist.to_dask_array(flow=False, dd=False)  # type: ignore
-    return hist
+    return hist  # type: ignore
 
 
 def histogram(
@@ -794,4 +794,4 @@ def histogram(
 
     if histogram != Histogram:
         return hist.to_dask_array(flow=False, dd=False)  # type: ignore
-    return hist
+    return hist  # type: ignore

src/dask_histogram/core.py

@@ -253,7 +253,7 @@ def _rebuild(
         name = self._name
         if rename:
             name = rename.get(name, name)
-        return type(self)(dsk, name, self.meta)
+        return type(self)(dsk, name, self.histref)
 
     @property
     def name(self) -> str:
@@ -312,7 +312,7 @@ def __setstate__(self, state: tuple[HighLevelGraph, str, bh.Histogram]) -> None:
 
     def to_dask_array(
         self, flow: bool = False, dd: bool = False
-    ) -> tuple[da.Array, ...] | tuple[da.Array, tuple[da.Array, ...]]:
+    ) -> tuple[da.Array, ...] | tuple[da.Array, list[da.Array]]:
         """Convert histogram object to dask.array form.
 
         Parameters
@@ -468,7 +468,7 @@ def __str__(self) -> str:
 
     @property
     def _args(self) -> tuple[HighLevelGraph, str, int, bh.Histogram]:
-        return (self.dask, self.name, self.npartititions, self.histref)
+        return (self.dask, self.name, self.npartitions, self.histref)
 
     def __getstate__(self) -> tuple[HighLevelGraph, str, int, bh.Histogram]:
         return self._args
@@ -564,7 +564,7 @@ def _partitioned_histogram(
     weights: DaskCollection | None = None,
     sample: DaskCollection | None = None,
     split_every: int | None = None,
-) -> AggHistogram:
+) -> PartitionedHistogram:
     name = f"hist-on-block-{tokenize(data, histref, weights, sample)}"
     data_is_df = is_dataframe_like(data[0])
     _weight_sample_check(*data, weights=weights)
@@ -630,7 +630,7 @@ def to_dask_array(
     agghist: AggHistogram,
     flow: bool = False,
     dd: bool = False,
-) -> tuple[DaskCollection, list[DaskCollection]] | tuple[DaskCollection, ...]:
+) -> tuple[da.Array, ...] | tuple[da.Array, list[da.Array]]:
     """Convert `agghist` to a `dask.array` return style.
 
     Parameters
@@ -694,12 +694,12 @@ def __call__(self, a: AggHistogram, b: AggHistogram) -> AggHistogram:
             deps.append(a)
             k1 = a.name
         else:
-            k1 = a
+            k1 = a  # type: ignore
         if is_dask_collection(b):
             deps.append(b)
             k2 = b.name
         else:
-            k2 = b
+            k2 = b  # type: ignore
         k1 = a.__dask_tokenize__() if is_dask_collection(a) else a  # type: ignore
         k2 = b.__dask_tokenize__() if is_dask_collection(b) else b  # type: ignore
         llg = {name: (self.func, k1, k2)}
@@ -832,7 +832,7 @@ def factory(
             storage = bh.storage.Double()
         histref = bh.Histogram(*axes, storage=storage)  # type: ignore
     f = _partitioned_histogram if keep_partitioned else _reduced_histogram
-    return f(
+    return f(  # type: ignore
         *data,
         histref=histref,
         weights=weights,

src/dask_histogram/routines.py

@@ -32,7 +32,7 @@ def histogram(
     histogram: Any | None = None,
     storage: bh.storage.Storage = bh.storage.Double(),
     threads: int | None = None,
-) -> AggHistogram | tuple[da.Array, ...]:
+) -> AggHistogram | tuple[da.Array, ...] | tuple[da.Array, list[da.Array]]:
     """Histogram Dask data in one dimension.
 
     Parameters
@@ -118,7 +118,7 @@ def histogram(
     )
     if histogram is None:
         return h.to_dask_array(flow=False, dd=False)  # type: ignore
-    return h
+    return h  # type: ignore
 
 
 def histogram2d(
@@ -241,7 +241,7 @@ def histogram2d(
     )
     if histogram is None:
         return h.to_dask_array(flow=False, dd=False)  # type: ignore
-    return h
+    return h  # type: ignore
 
 
 def histogramdd(
@@ -255,7 +255,7 @@ def histogramdd(
     histogram: Any | None = None,
     storage: bh.storage.Storage = bh.storage.Double(),
     threads: int | None = None,
-) -> (AggHistogram | tuple[da.Array, ...] | tuple[da.Array, tuple[da.Array, ...]]):
+) -> AggHistogram | tuple[da.Array, list[da.Array]]:
     """Histogram Dask data in multiple dimensions.
 
     Parameters
@@ -443,5 +443,5 @@ def histogramdd(
     ah = factory(*a, axes=axes, storage=storage, weights=weights)
 
     if histogram is not None:
-        return ah
-    return ah.to_dask_array(flow=False, dd=True)
+        return ah  # type: ignore
+    return ah.to_dask_array(flow=False, dd=True)  # type: ignore

tests/test_core.py

@@ -4,9 +4,9 @@
 import dask.array as da
 import dask.array.utils as dau
 import dask.datasets as dds
-import dask.delayed as delayed
 import numpy as np
 import pytest
+from dask.delayed import delayed
 
 import dask_histogram.core as dhc
 
@@ -152,7 +152,7 @@ def gen_hist_1D(
     range: tuple[float, float] = (-3, 3),
     size: tuple[int, ...] = (1000,),
     chunks: tuple[int, ...] = (250,),
-) -> dhc.AggHistogram:
+):
     hr = bh.Histogram(
         bh.axis.Regular(bins, range[0], range[1]),
         storage=bh.storage.Weight(),