Merge pull request #20063 from BvB93/histograms

ENH: Add annotations for `np.lib.histograms`

numpy/lib/histograms.pyi

@@ -1,7 +1,51 @@
-from typing import List
+from typing import (
+    Literal as L,
+    List,
+    Tuple,
+    Any,
+    SupportsIndex,
+    Sequence,
+)
+
+from numpy.typing import (
+    NDArray,
+    ArrayLike,
+)
+
+_BinKind = L[
+    "stone",
+    "auto",
+    "doane",
+    "fd",
+    "rice",
+    "scott",
+    "sqrt",
+    "sturges",
+]
 
 __all__: List[str]
 
-def histogram_bin_edges(a, bins=..., range=..., weights=...): ...
-def histogram(a, bins=..., range=..., normed=..., weights=..., density=...): ...
-def histogramdd(sample, bins=..., range=..., normed=..., weights=..., density=...): ...
+def histogram_bin_edges(
+    a: ArrayLike,
+    bins: _BinKind | SupportsIndex | ArrayLike = ...,
+    range: None | Tuple[float, float] = ...,
+    weights: None | ArrayLike = ...,
+) -> NDArray[Any]: ...
+
+def histogram(
+    a: ArrayLike,
+    bins: _BinKind | SupportsIndex | ArrayLike = ...,
+    range: None | Tuple[float, float] = ...,
+    normed: None = ...,
+    weights: None | ArrayLike = ...,
+    density: bool = ...,
+) -> Tuple[NDArray[Any], NDArray[Any]]: ...
+
+def histogramdd(
+    sample: ArrayLike,
+    bins: SupportsIndex | ArrayLike = ...,
+    range: Sequence[Tuple[float, float]] = ...,
+    normed: None | bool = ...,
+    weights: None | ArrayLike = ...,
+    density: None | bool = ...,
+) -> Tuple[NDArray[Any], List[NDArray[Any]]]: ...

numpy/typing/tests/data/fail/histograms.pyi

@@ -0,0 +1,13 @@
+import numpy as np
+import numpy.typing as npt
+
+AR_i8: npt.NDArray[np.int64]
+AR_f8: npt.NDArray[np.float64]
+
+np.histogram_bin_edges(AR_i8, range=(0, 1, 2))  # E: incompatible type
+
+np.histogram(AR_i8, range=(0, 1, 2))  # E: incompatible type
+np.histogram(AR_i8, normed=True)  # E: incompatible type
+
+np.histogramdd(AR_i8, range=(0, 1))  # E: incompatible type
+np.histogramdd(AR_i8, range=[(0, 1, 2)])  # E: incompatible type

numpy/typing/tests/data/reveal/histograms.pyi

@@ -0,0 +1,19 @@
+import numpy as np
+import numpy.typing as npt
+
+AR_i8: npt.NDArray[np.int64]
+AR_f8: npt.NDArray[np.float64]
+
+reveal_type(np.histogram_bin_edges(AR_i8, bins="auto"))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.histogram_bin_edges(AR_i8, bins="rice", range=(0, 3)))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.histogram_bin_edges(AR_i8, bins="scott", weights=AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.histogram(AR_i8, bins="auto"))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[Any]], numpy.ndarray[Any, numpy.dtype[Any]]]
+reveal_type(np.histogram(AR_i8, bins="rice", range=(0, 3)))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[Any]], numpy.ndarray[Any, numpy.dtype[Any]]]
+reveal_type(np.histogram(AR_i8, bins="scott", weights=AR_f8))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[Any]], numpy.ndarray[Any, numpy.dtype[Any]]]
+reveal_type(np.histogram(AR_f8, bins=1, density=True))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[Any]], numpy.ndarray[Any, numpy.dtype[Any]]]
+
+reveal_type(np.histogramdd(AR_i8, bins=[1]))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[Any]], builtins.list[numpy.ndarray[Any, numpy.dtype[Any]]]]
+reveal_type(np.histogramdd(AR_i8, range=[(0, 3)]))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[Any]], builtins.list[numpy.ndarray[Any, numpy.dtype[Any]]]]
+reveal_type(np.histogramdd(AR_i8, weights=AR_f8))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[Any]], builtins.list[numpy.ndarray[Any, numpy.dtype[Any]]]]
+reveal_type(np.histogramdd(AR_f8, density=True))  # E: Tuple[numpy.ndarray[Any, numpy.dtype[Any]], builtins.list[numpy.ndarray[Any, numpy.dtype[Any]]]]