TYP: bootstrap type-checking

pyproject.toml

@@ -115,3 +115,12 @@ omit = [
     "unyt/_on_demand_imports.py",
     "unyt/tests/test_linters.py",
 ]
+
+[tool.mypy]
+python_version = 3.8
+show_error_codes = true
+ignore_missing_imports = true
+warn_unused_configs = true
+warn_unused_ignores = true
+warn_unreachable = true
+show_error_context = true

tox.ini

@@ -1,10 +1,10 @@
 [tox]
-envlist = py38-docs,begin,py38-dependencies,py38-versions,py{38,39,310},py38-unyt-module-test-function,end
+envlist = py38-docs,begin,py38-dependencies,py38-versions,py{38,39,310},py38-unyt-module-test-function,typecheck,end
 isolated_build = True
 
 [gh-actions]
 python =
-    3.8: py38, py38-docs, py38-dependencies, py38-versions, py38-unyt-module-test-function
+    3.8: py38, py38-docs, py38-dependencies, py38-versions, py38-unyt-module-test-function, typecheck
     3.9: py39
     3.10: py310
 
@@ -79,6 +79,15 @@ depends = py38
 commands =
     python -c 'import unyt; unyt.test()'
 
+[testenv:typecheck]
+deps =
+    mypy==0.991
+    numpy==1.17.5
+    sympy==1.5
+
+commands =
+    python -m mypy unyt
+
 [testenv:begin]
 commands =
     coverage erase

unyt/__init__.py

@@ -93,6 +93,6 @@ def test():  # pragma: no cover
 
 
 # isort: off
-from unyt.mpl_interface import matplotlib_support
+from unyt.mpl_interface import MatplotlibSupport
 
-matplotlib_support = matplotlib_support()
+matplotlib_support = MatplotlibSupport()

unyt/_array_functions.py

@@ -4,7 +4,7 @@
 import numpy as np
 from packaging.version import Version
 
-from unyt import delta_degC
+from unyt import delta_degC  # type: ignore [attr-defined]
 from unyt.array import NULL_UNIT, unyt_array
 from unyt.dimensions import temperature
 from unyt.exceptions import (

unyt/_mpl_array_converter/__init__.py

@@ -1,17 +1,28 @@
+from typing import TYPE_CHECKING
 from weakref import WeakKeyDictionary
 
 from matplotlib.units import AxisInfo, ConversionInterface
 
 from unyt.array import unyt_array, unyt_quantity
 from unyt.unit_object import Unit
 
+if TYPE_CHECKING:
+    import sys
+
+    from matplotlib.axes import Axis
+
+    if sys.version_info >= (3, 9):
+        from collections.abc import MutableMapping
+    else:
+        from typing import MutableMapping
+
 
 class unyt_arrayConverter(ConversionInterface):
     """Matplotlib interface for unyt_array"""
 
     _instance = None
     _labelstyle = "()"
-    _axisnames = WeakKeyDictionary()
+    _axisnames: "MutableMapping[Axis, str]" = WeakKeyDictionary()
 
     # ensure that unyt_arrayConverter is a singleton
     def __new__(cls):

unyt/array.py

@@ -28,6 +28,7 @@
     bitwise_or,
     bitwise_xor,
     ceil,
+    clip,
     conj,
     copysign,
     cos,
@@ -101,7 +102,7 @@
     true_divide,
     trunc,
 )
-from numpy.core.umath import _ones_like, clip
+from numpy.core.umath import _ones_like
 from sympy import Rational
 
 from unyt._on_demand_imports import _astropy, _dask, _pint

unyt/mpl_interface.py

@@ -13,10 +13,10 @@
 
 from ._on_demand_imports import _matplotlib
 
-__all__ = ["matplotlib_support"]
+__all__ = ["MatplotlibSupport"]
 
 
-class matplotlib_support:
+class MatplotlibSupport:
     """Context manager for enabling the feature
 
     When used in a with statement, the feature is enabled during the context and

unyt/tests/test_array_functions.py

@@ -6,7 +6,7 @@
 import pytest
 from packaging.version import Version
 
-from unyt import K, cm, degC, degF, delta_degC, g, km, s
+from unyt import K, cm, degC, degF, delta_degC, g, km, s  # type: ignore [attr-defined]
 from unyt._array_functions import _HANDLED_FUNCTIONS as HANDLED_FUNCTIONS
 from unyt.array import unyt_array, unyt_quantity
 from unyt.exceptions import (

unyt/tests/test_dask_arrays.py

@@ -1,5 +1,6 @@
 import pickle
 from collections import defaultdict
+from typing import Any, DefaultDict, Tuple
 
 import numpy as np
 import pytest
@@ -19,7 +20,7 @@
     unyt_from_dask,
 )
 from unyt.exceptions import UnitOperationError
-from unyt.unit_symbols import cm, g, m
+from unyt.unit_symbols import cm, g, m  # type: ignore [attr-defined]
 
 
 def test_unyt_dask_creation():
@@ -263,7 +264,7 @@ def test_unyt_type_result():
     assert result == unyt_quantity(1, m)
 
 
-_func_args = defaultdict(lambda: ())
+_func_args: DefaultDict[str, Tuple[Any, ...]] = defaultdict(lambda: ())
 _func_args["reshape"] = ((100, 1),)
 _func_args["rechunk"] = ((5, 5),)
 _func_args["cumsum"] = (0,)

unyt/tests/test_mpl_interface.py

@@ -2,7 +2,14 @@
 import numpy as np
 import pytest
 
-from unyt import K, m, matplotlib_support, s, unyt_array, unyt_quantity
+from unyt import (  # type: ignore [attr-defined]
+    K,
+    m,
+    matplotlib_support,
+    s,
+    unyt_array,
+    unyt_quantity,
+)
 from unyt._on_demand_imports import NotAModule, _matplotlib
 
 try:

unyt/tests/test_unyt_array.py

@@ -29,7 +29,16 @@
 )
 from packaging.version import Version
 
-from unyt import K, R, Unit, degC, degF, delta_degC, delta_degF, dimensions
+from unyt import (  # type: ignore [attr-defined]
+    K,
+    R,
+    Unit,
+    degC,
+    degF,
+    delta_degC,
+    delta_degF,
+    dimensions,
+)
 from unyt._on_demand_imports import _astropy, _h5py, _pint
 from unyt._physical_ratios import metallicity_sun, speed_of_light_cm_per_s
 from unyt.array import (
@@ -63,7 +72,7 @@
     assert_array_equal_units,
 )
 from unyt.unit_registry import UnitRegistry
-from unyt.unit_symbols import cm, degree, g, m
+from unyt.unit_symbols import cm, degree, g, m  # type: ignore [attr-defined]
 
 
 def operate_and_compare(a, b, op, answer):