WIP

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

Author: nstarman

.github/workflows/ci_workflows.yml

@@ -3,9 +3,9 @@ name: CI
 on:
   push:
     branches:
-    - main
+      - main
     tags:
-    - '*'
+      - "*"
   pull_request:
 
 concurrency:

.pre-commit-config.yaml

@@ -0,0 +1,89 @@
+ci:
+  autoupdate_schedule: "monthly"
+  autoupdate_commit_msg: "chore: update pre-commit hooks"
+  autofix_commit_msg: "style: pre-commit fixes"
+
+default_stages: [pre-commit, pre-push]
+
+repos:
+  - repo: meta
+    hooks:
+      - id: check-useless-excludes
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: "v5.0.0"
+    hooks:
+      - id: check-added-large-files
+      - id: check-case-conflict
+      - id: check-merge-conflict
+      - id: check-yaml
+      - id: debug-statements
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+      - id: name-tests-test
+        args: ["--pytest-test-first"]
+      - id: trailing-whitespace
+
+  - repo: https://github.com/pre-commit/pygrep-hooks
+    rev: "v1.10.0"
+    hooks:
+      - id: rst-backticks
+      - id: rst-directive-colons
+      - id: rst-inline-touching-normal
+
+  - repo: https://github.com/python-jsonschema/check-jsonschema
+    rev: 0.29.4
+    hooks:
+      - id: check-dependabot
+      - id: check-github-workflows
+      - id: check-readthedocs
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: "v0.7.3"
+    hooks:
+      # Run the linter
+      - id: ruff
+        types_or: [python, pyi, jupyter]
+        args: ["--fix", "--show-fixes"]
+      # Run the formatter
+      - id: ruff-format
+        types_or: [python, pyi, jupyter]
+
+  - repo: https://github.com/adamchainz/blacken-docs
+    rev: "1.19.1"
+    hooks:
+      - id: blacken-docs
+        additional_dependencies: [black==23.*]
+
+  - repo: https://github.com/rbubley/mirrors-prettier
+    rev: "v3.3.3"
+    hooks:
+      - id: prettier
+        types_or: [yaml, markdown, html, css, scss, javascript, json]
+        args: [--prose-wrap=always]
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: "v1.13.0"
+    hooks:
+      - id: mypy
+        files: src
+        additional_dependencies:
+          - pytest
+
+  - repo: https://github.com/codespell-project/codespell
+    rev: "v2.3.0"
+    hooks:
+      - id: codespell
+
+  - repo: https://github.com/abravalheri/validate-pyproject
+    rev: v0.23
+    hooks:
+      - id: validate-pyproject
+
+  - repo: local
+    hooks:
+      - id: disallow-caps
+        name: Disallow improper capitalization
+        language: pygrep
+        entry: PyBind|Numpy|Cmake|CCache|Github|PyTest
+        exclude: .pre-commit-config.yaml

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
   "array-api-compat>=1.9.1",
   "astropy>=7.0",
   "numpy>=2.0",
+  "typing-extensions>=4.12.2",
 ]
 dynamic = ["version"]
 
@@ -103,6 +104,30 @@ exclude_lines = [
   "@overload",
 ]
 
+[tool.mypy]
+  python_version           = "3.11"
+  files                    = ["quantity"]
+  strict                   = true
+
+  disallow_incomplete_defs = true
+  disallow_untyped_defs    = false
+  enable_error_code        = ["ignore-without-code", "redundant-expr", "truthy-bool"]
+  warn_return_any          = true
+  warn_unreachable         = true
+  warn_unused_configs      = true
+
+  [[tool.mypy.overrides]]
+    module = ["quantity._dev.*", "quantity.tests.*"]
+    ignore_errors = true
+
+  [[tool.mypy.overrides]]
+    ignore_missing_imports = true
+    module = [
+      "astropy.*",
+      "array_api_compat.*"
+    ]
+
+
 [tool.ruff]
 exclude=[  # package template provided files.
     "setup.py",
@@ -140,10 +165,12 @@ ignore = [
   "PLR2004",  # Magic value used in comparison
   "RET505",   # Unnecessary `else`/`elif` after `return` statement
 ]
-isort.required-imports = ["from __future__ import annotations"]
-# Uncomment if using a _compat.typing backport
-# typing-modules = ["quantity_2_0._compat.typing"]
 
 [tool.ruff.lint.per-file-ignores]
 "tests/**" = ["T20"]
 "noxfile.py" = ["T20"]
+
+[dependency-groups]
+typing = [
+    "mypy>=1.13.0",
+]

quantity/_array_api.py

@@ -0,0 +1,17 @@
+"""Minimal definition of the Array API."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+
+class HasArrayNameSpace(Protocol):
+    """Minimal defintion of the Array API."""
+
+    def __array_namespace__(self) -> Any: ...
+
+
+class Array(HasArrayNameSpace, Protocol):
+    """Minimal defintion of the Array API."""
+
+    def __pow__(self, other: Any) -> Array: ...

quantity/_quantity_api.py

@@ -0,0 +1,24 @@
+"""Minimal definition of the Quantity API."""
+
+__all__ = ["Quantity", "ArrayQuantity", "Unit"]
+
+from typing import Protocol, runtime_checkable
+
+from astropy.units import UnitBase as Unit
+
+from ._array_api import Array
+
+
+@runtime_checkable
+class Quantity(Protocol):
+    """Minimal definition of the Quantity API."""
+
+    value: Array
+    unit: Unit
+
+
+@runtime_checkable
+class ArrayQuantity(Quantity, Array, Protocol):
+    """An array-valued Quantity."""
+
+    ...

quantity/core.py

@@ -1,46 +1,63 @@
+"""Quantity."""
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
+
 from __future__ import annotations
 
 import operator
+from collections.abc import Callable
 from dataclasses import dataclass, replace
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    TypeAlias,
+    TypeGuard,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 
 import array_api_compat
 import astropy.units as u
 import numpy as np
+from astropy.units import UnitBase as Unit
 from astropy.units.quantity_helper import UFUNC_HELPERS
 
 if TYPE_CHECKING:
-    from typing import Any
+    from types import NotImplementedType
+    from typing import Any, Self
+
+    from ._array_api import Array
+    from ._quantity_api import ArrayQuantity, Unit
+
+
+T = TypeVar("T")
 
 
 DIMENSIONLESS = u.dimensionless_unscaled
 
 PYTHON_NUMBER = float | int | complex
 
 
-def has_array_namespace(arg):
+def has_array_namespace(arg: Any) -> TypeGuard[Array]:
     try:
         array_api_compat.array_namespace(arg)
     except TypeError:
         return False
-    else:
-        return True
+    return True
 
 
-def get_value_and_unit(arg, default_unit=None):
-    # HACK: interoperability with astropy Quantity.  Have protocol?
-    try:
-        unit = arg.unit
-    except AttributeError:
-        return arg, default_unit
-    else:
-        return arg.value, unit
+def get_value_and_unit(
+    arg: ArrayQuantity | Array, default_unit: Unit | None = None
+) -> tuple[Array, Unit]:
+    return (
+        (arg.value, arg.unit) if isinstance(arg, ArrayQuantity) else (arg, default_unit)
+    )
 
 
-def value_in_unit(value, unit):
+def value_in_unit(value: Array, unit: Unit) -> Array:
     v_value, v_unit = get_value_and_unit(value, default_unit=DIMENSIONLESS)
-    return v_unit.to(unit, v_value)
+    return cast(Array, v_unit.to(unit, v_value))
 
 
 _OP_TO_NP_FUNC = {
@@ -55,7 +72,12 @@ def value_in_unit(value, unit):
 OP_HELPERS = {op: UFUNC_HELPERS[np_func] for op, np_func in _OP_TO_NP_FUNC.items()}
 
 
-def _make_op(fop, mode):
+QuantityOpCallable: TypeAlias = Callable[
+    ["Quantity", Any], Union["Quantity", NotImplementedType]
+]
+
+
+def _make_op(fop: str, mode: str) -> QuantityOpCallable:
     assert mode in "fri"
     op = fop if mode == "f" else "__" + mode + fop[2:]
     helper = OP_HELPERS[fop]
@@ -75,27 +97,29 @@ def __op__(self, other):
     return __op__
 
 
-def _make_ops(op):
-    return tuple(_make_op(op, mode) for mode in "fri")
+def _make_ops(
+    op: str,
+) -> tuple[QuantityOpCallable, QuantityOpCallable, QuantityOpCallable]:
+    return (_make_op(op, "f"), _make_op(op, "r"), _make_op(op, "i"))
 
 
-def _make_comp(comp):
-    def __comp__(self, other):
+def _make_comp(comp: str) -> Callable[[Quantity, Any], Array]:
+    def _comp_(self: Quantity, other: Any) -> Array | NotImplementedType:
         try:
             other = value_in_unit(other, self.unit)
         except Exception:
             return NotImplemented
         return getattr(self.value, comp)(other)
 
-    return __comp__
+    return _comp_
 
 
-def _make_deferred(attr):
+def _make_deferred(attr: str) -> Callable[[Quantity], property]:
     # Use array_api_compat getter if available (size, device), since
     # some array formats provide inconsistent implementations.
     attr_getter = getattr(array_api_compat, attr, operator.attrgetter(attr))
 
-    def deferred(self):
+    def deferred(self: Quantity):
         return attr_getter(self.value)
 
     return property(deferred)
@@ -133,33 +157,61 @@ def defer_dimensionless(self):
 
     return defer_dimensionless
 
+# -----------------
+
+
+@overload
+def _parse_pow_mod(mod: None, /) -> None: ...
+
+
+@overload
+def _parse_pow_mod(mod: object, /) -> NotImplementedType: ...
+
 
-def _check_pow_args(exp, mod):
-    if mod is not None:
-        return NotImplemented
+def _parse_pow_mod(mod: T, /) -> T | NotImplementedType:
+    return mod if mod is None else NotImplemented  # type: ignore[redundant-expr]
 
-    if not isinstance(exp, PYTHON_NUMBER):
+
+# -----------------
+
+
+@overload
+def _check_pow_exp(exp: Array | PYTHON_NUMBER, /) -> PYTHON_NUMBER: ...
+
+
+@overload
+def _check_pow_exp(exp: object, /) -> NotImplementedType: ...
+
+
+def _check_pow_exp(exp: Any, /) -> PYTHON_NUMBER | NotImplementedType:
+    out: PYTHON_NUMBER
+    if isinstance(exp, PYTHON_NUMBER):
+        out = exp
+    else:
         try:
-            exp = exp.__complex__()
+            out = complex(exp)
         except Exception:
             try:
-                return exp.__float__()
+                return float(exp)
             except Exception:
                 return NotImplemented
 
-    return exp.real if exp.imag == 0 else exp
+    return out.real if out.imag == 0 else out
 
 
 @dataclass(frozen=True, eq=False)
 class Quantity:
-    value: Any
-    unit: u.UnitBase
+    value: Array
+    unit: Unit
 
     def __array_namespace__(self, *, api_version: str | None = None) -> Any:
         # TODO: make our own?
+        del api_version
         return np
 
-    def _operate(self, other, op, units_helper):
+    def _operate(
+        self, other: Any, op: Any, units_helper: Any
+    ) -> Self | NotImplementedType:
         if not has_array_namespace(other) and not isinstance(other, PYTHON_NUMBER):
             # HACK: unit should take care of this!
             if not isinstance(other, u.UnitBase):
@@ -205,7 +257,7 @@ def _operate(self, other, op, units_helper):
     __lt__ = _make_comp("__lt__")
     __ne__ = _make_comp("__ne__")
 
-    # Atttributes deferred to those of .value
+    # Attributes deferred to those of .value
     dtype = _make_deferred("dtype")
     device = _make_deferred("device")
     ndim = _make_deferred("ndim")
@@ -228,9 +280,11 @@ def _operate(self, other, op, units_helper):
 
     # TODO: __dlpack__, __dlpack_device__
 
-    def __pow__(self, exp, mod=None):
-        exp = _check_pow_args(exp, mod)
-        if exp is NotImplemented:
+    def __pow__(self, exp: Any, mod: Any = None) -> Self | NotImplementedType:
+        if (mod := _parse_pow_mod(mod)) is NotImplemented:
+            return NotImplemented
+
+        if (exp := _check_pow_exp(exp)) is NotImplemented:
             return NotImplemented
 
         value = self.value.__pow__(exp)
@@ -239,17 +293,25 @@ def __pow__(self, exp, mod=None):
         return replace(self, value=value, unit=self.unit**exp)
 
     def __ipow__(self, exp, mod=None):
-        exp = _check_pow_args(exp, mod)
-        if exp is NotImplemented:
+        if (mod := _parse_pow_mod(mod)) is NotImplemented:
+            return NotImplemented
+
+        if (exp := _check_pow_exp(exp)) is NotImplemented:
             return NotImplemented
 
         value = self.value.__ipow__(exp)
         if value is NotImplemented:
             return NotImplemented
         return replace(self, value=value, unit=self.unit**exp)
 
-    def __setitem__(self, item, value):
-        self.value[item] = value_in_unit(value, self.unit)
+    def __setitem__(self, item: Any, value: Any) -> None:
+        """Call the setitem method of the array for the value in the unit.
+
+        The Array API does not guarantee mutability of the underlying array,
+        so this method will raise an exception if the array is immutable.
+
+        """
+        self.value[item] = value_in_unit(value, self.unit)  # type: ignore[index]
 
     __array_ufunc__ = None
     __array_function__ = None

quantity/py.typed

@@ -1,6 +1,9 @@
 # NOTE: First try _dev.scm_version if it exists and setuptools_scm is installed
 # This file is not included in wheels/tarballs, so otherwise it will
 # fall back on the generated _version module.
+
+__all__ = ['version']
+
 try:
     try:
         from ._dev.scm_version import version