chore: run mypy over subset of awkward

.pre-commit-config.yaml

@@ -66,3 +66,10 @@ repos:
   hooks:
     - id: check-github-workflows
       args: ["--verbose"]
+
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v1.6.1
+  hooks:
+    - id: mypy
+      additional_dependencies:
+        - numpy>=1.24

pyproject.toml

@@ -220,40 +220,21 @@ disable = [
 
 [tool.mypy]
 files = ["src/awkward/**/*.py"]
-exclude = ["^src/awkward/[^/]+\\.py$"]
 plugins = [
     "numpy.typing.mypy_plugin"
 ]
 python_version = "3.11"
+ignore_errors = true
+ignore_missing_imports = true
 
 [[tool.mypy.overrides]]
 module = [
-    'awkward.__init__',
-    'awkward._connect.*',
-    'awkward._cpu_kernels',
-    'awkward._errors',
-    'awkward._kernel_signatures',
-    'awkward._libawkward',
-    'awkward._util',
-    'awkward.forms',
-    'awkward.forth',
-    'awkward.highlevel',
-    'awkward.nplike',
-    'awkward.numba',
-    'awkward.types',
-    'awkward.types._awkward_datashape_parser',
-    'numba.*',
-    'llvmlite.*',
-    'ROOT.*',
-    'cppyy.*',
-    'jax.*',
-    'pandas.*',
-    'cupy.*',
-    'pyarrow.*',
-    'fsspec.*',
-    'numexpr.*',
+    'awkward._nplikes.*',
+    'awkward._behavior.*',
+    'awkward._backends.*',
+    'awkward.forms.*',
 ]
-ignore_errors = true
+ignore_errors = false
 ignore_missing_imports = true
 
 [tool.ruff]
@@ -306,6 +287,7 @@ mccabe.max-complexity = 100
 [tool.ruff.lint.per-file-ignores]
 "dev/*" = ["T20", "TID251"]
 "src/awkward/numba/*" = ["TID251"]
+"src/awkward/_typing.py" = ["TID251"]
 "src/awkward/_errors.py" = ["TID251"]
 "src/awkward/errors.py" = ["TID251"]
 "src/awkward/_connect/*" = ["TID251"]

src/awkward/_backends/backend.py

@@ -2,20 +2,21 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from typing import Any
 
 import awkward as ak
 from awkward._kernels import KernelError
 from awkward._nplikes.numpy import Numpy
 from awkward._nplikes.numpylike import NumpyLike, NumpyMetadata
 from awkward._singleton import PublicSingleton
-from awkward._typing import Callable, Tuple, TypeAlias, TypeVar, Unpack
+from awkward._typing import Callable, Tuple, TypeAlias, TypeVar
 
 np = NumpyMetadata.instance()
 numpy = Numpy.instance()
 
 
 T_co = TypeVar("T_co", covariant=True)
-KernelKeyType: TypeAlias = Tuple[str, Unpack[Tuple[np.dtype, ...]]]
+KernelKeyType: TypeAlias = Tuple[Any, ...]
 KernelType: TypeAlias = "Callable[..., KernelError | None]"
 
 
@@ -49,6 +50,7 @@ def format_kernel_error(
                 errors="surrogateescape"
             ).lstrip("\n").lstrip("(")
 
+        assert error.str is not None
         message = error.str.decode(errors="surrogateescape")
 
         if error.attempt != ak._util.kSliceNone:

src/awkward/_backends/cupy.py

@@ -13,7 +13,7 @@
 numpy = Numpy.instance()
 
 
-@register_backend(Cupy)
+@register_backend(Cupy)  # type: ignore[type-abstract]
 class CupyBackend(Backend):
     name: Final[str] = "cuda"
 

src/awkward/_backends/dispatch.py

@@ -6,8 +6,8 @@
 from awkward._backends.backend import Backend
 from awkward._nplikes.numpy import Numpy
 from awkward._nplikes.numpylike import NumpyLike, NumpyMetadata
-from awkward._typing import Callable, TypeAlias, TypeVar
-from awkward._util import UNSET
+from awkward._typing import Callable, TypeAlias, TypeVar, cast
+from awkward._util import UNSET, Sentinel
 
 np = NumpyMetadata.instance()
 numpy = Numpy.instance()
@@ -19,7 +19,7 @@
 BackendLookupFactory: TypeAlias = "Callable[[type[T]], BackendLookup[T]]"
 
 
-_type_to_backend_lookup: dict[type[T], BackendLookup] = {}
+_type_to_backend_lookup: dict[type, BackendLookup] = {}
 _backend_lookup_factories: list[BackendLookupFactory] = []
 _name_to_backend_cls: dict[str, type[Backend]] = {}
 
@@ -68,7 +68,7 @@ def common_backend(backends: Collection[Backend]) -> Backend:
             )
 
 
-def backend_of_obj(obj, default: D = UNSET) -> Backend | D:
+def backend_of_obj(obj, default: D | Sentinel = UNSET) -> Backend | D:
     cls = type(obj)
     try:
         lookup = _type_to_backend_lookup[cls]
@@ -82,13 +82,13 @@ def backend_of_obj(obj, default: D = UNSET) -> Backend | D:
             if default is UNSET:
                 raise TypeError(f"cannot find backend for {cls.__name__}")
             else:
-                return default
+                return cast(D, default)
         _type_to_backend_lookup[cls] = maybe_lookup
         return maybe_lookup(obj)
 
 
 def backend_of(
-    *objects, default: D = UNSET, coerce_to_common: bool = False
+    *objects, default: D | Sentinel = UNSET, coerce_to_common: bool = False
 ) -> Backend | D:
     """
     Args:
@@ -108,7 +108,7 @@ def backend_of(
         if default is UNSET:
             raise ValueError("could not find backend for", objects)
         else:
-            return default
+            return cast(D, default)
     elif len(unique_backends) == 1:
         return next(iter(unique_backends))
     elif coerce_to_common:

src/awkward/_backends/jax.py

@@ -16,7 +16,7 @@
 numpy = Numpy.instance()
 
 
-@register_backend(Jax)
+@register_backend(Jax)  # type: ignore[type-abstract]
 class JaxBackend(Backend):
     name: Final[str] = "jax"
 

src/awkward/_backends/numpy.py

@@ -14,7 +14,7 @@
 numpy = Numpy.instance()
 
 
-@register_backend(Numpy)
+@register_backend(Numpy)  # type: ignore[type-abstract]
 class NumpyBackend(Backend):
     name: Final[str] = "cpu"
 

src/awkward/_backends/typetracer.py

@@ -13,7 +13,7 @@
 numpy = Numpy.instance()
 
 
-@register_backend(TypeTracer)
+@register_backend(TypeTracer)  # type: ignore[type-abstract]
 class TypeTracerBackend(Backend):
     name: Final[str] = "typetracer"
 

src/awkward/_behavior.py

@@ -2,14 +2,22 @@
 from __future__ import annotations
 
 from collections import ChainMap
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 
 import awkward as ak
 from awkward._nplikes import ufuncs
-from awkward._typing import Any
+from awkward._typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from awkward._nplikes.numpylike import UfuncLike
+    from awkward._reducers import Reducer
+    from awkward.contents.content import Content
+    from awkward.highlevel import Array
+    from awkward.highlevel import Record as HighLevelRecord
+    from awkward.record import Record
 
 
-def overlay_behavior(behavior: dict | None) -> Mapping:
+def overlay_behavior(behavior: Mapping | None) -> Mapping:
     """
     Args:
         behavior: behavior dictionary, or None
@@ -19,10 +27,11 @@ def overlay_behavior(behavior: dict | None) -> Mapping:
     """
     if behavior is None:
         return ak.behavior
-    return ChainMap(behavior, ak.behavior)
+    else:
+        return ChainMap(behavior, ak.behavior)  # type: ignore[arg-type]
 
 
-def get_array_class(layout, behavior):
+def get_array_class(layout: Content, behavior: Mapping | None) -> type[Array]:
     from awkward.highlevel import Array
 
     behavior = overlay_behavior(behavior)
@@ -39,7 +48,7 @@ def get_array_class(layout, behavior):
     return Array
 
 
-def get_record_class(layout, behavior):
+def get_record_class(layout: Record, behavior: Mapping | None) -> type[HighLevelRecord]:
     from awkward.highlevel import Record
 
     behavior = overlay_behavior(behavior)
@@ -51,14 +60,20 @@ def get_record_class(layout, behavior):
     return Record
 
 
-def find_record_reducer(reducer, layout, behavior):
+def find_record_reducer(
+    reducer: Reducer, layout: Record, behavior: Mapping | None
+) -> Callable[[Array, bool], Any] | None:
     behavior = overlay_behavior(behavior)
     rec = layout.parameter("__record__")
     if isinstance(rec, str):
         return behavior.get((reducer.highlevel_function(), rec))
+    else:
+        return None
 
 
-def find_custom_cast(obj, behavior):
+def find_custom_cast(
+    obj: Any, behavior: Mapping | None
+) -> Callable[[Any], Content | Record] | None:
     behavior = overlay_behavior(behavior)
     for cls in type(obj).__mro__:
         fcn = behavior.get(("__cast__", cls))
@@ -67,7 +82,9 @@ def find_custom_cast(obj, behavior):
     return None
 
 
-def find_ufunc_generic(ufunc, layout, behavior):
+def find_ufunc_generic(
+    ufunc: UfuncLike, layout: Content, behavior: Mapping | None
+) -> Callable[[UfuncLike, str, list, dict], Any] | None:
     behavior = overlay_behavior(behavior)
     custom = layout.parameter("__list__")
     if custom is None:
@@ -81,8 +98,8 @@ def find_ufunc_generic(ufunc, layout, behavior):
         return None
 
 
-def find_ufunc(behavior, signature: tuple):
-    if not any(s is None for s in signature):
+def find_ufunc(behavior: Mapping | None, signature: tuple) -> UfuncLike | None:
+    if all(s is not None for s in signature):
         behavior = overlay_behavior(behavior)
 
         # Special case all strings or hashable types.
@@ -105,13 +122,14 @@ def find_ufunc(behavior, signature: tuple):
                     )
                 ):
                     return custom
+    return None
 
 
 def find_record_typestr(
     behavior: None | Mapping,
     parameters: None | Mapping[str, Any],
     default: str | None = None,
-):
+) -> str | None:
     if parameters is None:
         return default
     behavior = overlay_behavior(behavior)
@@ -122,14 +140,14 @@ def find_array_typestr(
     behavior: None | Mapping,
     parameters: None | Mapping[str, Any],
     default: str | None = None,
-):
+) -> str | None:
     if parameters is None:
         return default
     behavior = overlay_behavior(behavior)
     return behavior.get(("__typestr__", parameters.get("__list__")), default)
 
 
-def behavior_of_obj(obj, behavior: Mapping | None = None) -> Mapping | None:
+def behavior_of_obj(obj: Any, behavior: Mapping | None = None) -> Mapping | None:
     from awkward.highlevel import Array, ArrayBuilder, Record
 
     if behavior is not None:
@@ -140,7 +158,7 @@ def behavior_of_obj(obj, behavior: Mapping | None = None) -> Mapping | None:
         return None
 
 
-def behavior_of(*arrays, behavior: Mapping | None = None) -> Mapping | None:
+def behavior_of(*arrays: Any, behavior: Mapping | None = None) -> Mapping | None:
     if behavior is not None:
         # An explicit 'behavior' always wins.
         return behavior
@@ -158,5 +176,6 @@ def behavior_of(*arrays, behavior: Mapping | None = None) -> Mapping | None:
             behavior.update(x_behavior)
             copied = True
         else:
+            assert isinstance(behavior, dict)
             behavior.update(x_behavior)
     return behavior

src/awkward/_kernels.py

@@ -21,8 +21,8 @@
 
 
 class KernelError(Protocol):
-    filename: str | None  # pylint: disable=E0602
-    str: str | None
+    filename: bytes | None  # pylint: disable=E0602
+    str: bytes | None
     attempt: int
     id: int
 

src/awkward/_nplikes/__init__.py

@@ -33,6 +33,6 @@ def to_nplike(
     if isinstance(from_nplike, awkward._nplikes.cupy.Cupy) and not isinstance(
         nplike, awkward._nplikes.cupy.Cupy
     ):
-        array = array.get()
+        array = array.get()  # type: ignore[attr-defined]
 
     return nplike.asarray(array)