Merge pull request #3546 from neutrinoceros/mypy-setup

TYP: setup basic type checking and fix existing type errors

.github/workflows/type-checking.yaml

@@ -0,0 +1,38 @@
+name: type checking
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    paths-ignore:
+      - "doc/**"
+      - README.md
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    name: mypy
+    timeout-minutes: 60
+
+    concurrency:
+      group: ${{ github.ref }}-dev
+      cancel-in-progress: true
+
+    steps:
+    - uses: actions/checkout@v2
+      with:
+        fetch-depth: 0
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+
+    - name: Build yt + type check deps
+      run: |
+        python3 -m pip install --upgrade pip
+        python3 -m pip install -e .[typecheck]
+
+    - name: Run mypy
+      run: mypy yt

pyproject.toml

@@ -129,3 +129,15 @@ ignore = [
   "yt/utilities/lib/ragged_arrays.c",
   "yt/utilities/lib/write_array.c",
 ]
+
+
+[tool.mypy]
+python_version = 3.7
+ignore_missing_imports = true
+warn_unused_configs = true
+warn_unused_ignores = true
+warn_unreachable = true
+show_error_context = true
+
+# color_maps can be reactivated after https://github.com/yt-project/yt/pull/3441 is merged
+exclude = "(test_*|lodgeit|color_maps)"

setup.cfg

@@ -107,6 +107,13 @@ test =
     nose-timer~=1.0.0
     pytest>=6.1
     pytest-xdist~=2.1.0
+typecheck =
+    mypy==0.910
+    types-PyYAML==5.4.10
+    types-chardet==4.0.0
+    types-requests==2.25.9
+    types-setuptools==57.4.0
+    types-toml==0.10.0
 
 [flake8]
 max-line-length = 88

yt/__init__.py

@@ -69,7 +69,7 @@
     uunion1d,
     uvstack,
 )
-from yt.units.unit_object import define_unit
+from yt.units.unit_object import define_unit  # type: ignore
 from yt.utilities.logger import set_log_level, ytLogger as mylog
 
 frontends = _frontend_container()
@@ -89,7 +89,7 @@
     load_unstructured_mesh,
 )
 from yt.testing import run_nose
-from yt.units.unit_systems import UnitSystem, unit_system_registry
+from yt.units.unit_systems import UnitSystem, unit_system_registry  # type: ignore
 
 # Import some helpful math utilities
 from yt.utilities.math_utils import ortho_find, periodic_position, quartiles

yt/_maintenance/backports.py

@@ -21,7 +21,7 @@
 import sys
 
 if sys.version_info < (3, 8):
-    from _thread import RLock
+    from _thread import RLock  # type: ignore
     from typing import List
 
     GenericAlias = type(List[int])

yt/_typing.py

@@ -0,0 +1,4 @@
+from typing import List, Optional, Tuple
+
+FieldDescT = Tuple[str, Tuple[str, List[str], Optional[str]]]
+KnownFieldsT = Tuple[FieldDescT, ...]

yt/config.py

@@ -190,7 +190,7 @@ def _repr_json_(self):
 _local_config_file = YTConfig.get_local_config_file()
 
 if not os.path.exists(_global_config_file):
-    cfg = {"yt": {}}
+    cfg = {"yt": {}}  # type: ignore
     try:
         with open(_global_config_file, mode="w") as fd:
             toml.dump(cfg, fd)

yt/data_objects/construction_data_containers.py

@@ -24,8 +24,8 @@
 from yt.geometry import particle_deposit as particle_deposit
 from yt.geometry.coordinates.cartesian_coordinates import all_data
 from yt.loaders import load_uniform_grid
-from yt.units.unit_object import Unit
-from yt.units.yt_array import YTArray, uconcatenate
+from yt.units.unit_object import Unit  # type: ignore
+from yt.units.yt_array import YTArray, uconcatenate  # type: ignore
 from yt.utilities.exceptions import (
     YTNoAPIKey,
     YTNotInsideNotebook,

yt/data_objects/data_containers.py

@@ -1,6 +1,8 @@
+import abc
 import weakref
 from collections import defaultdict
 from contextlib import contextmanager
+from typing import List, Tuple, Union
 
 import numpy as np
 
@@ -9,7 +11,7 @@
 from yt.fields.field_exceptions import NeedsGridType
 from yt.frontends.ytdata.utilities import save_as_dataset
 from yt.funcs import get_output_filename, is_sequence, iter_fields, mylog
-from yt.units.yt_array import YTArray, YTQuantity, uconcatenate
+from yt.units.yt_array import YTArray, YTQuantity, uconcatenate  # type: ignore
 from yt.utilities.amr_kdtree.api import AMRKDTree
 from yt.utilities.exceptions import (
     YTCouldNotGenerateField,
@@ -50,7 +52,7 @@ def _get_ipython_key_completion(ds):
     return tuple_keys + fnames
 
 
-class YTDataContainer:
+class YTDataContainer(abc.ABC):
     """
     Generic YTDataContainer container.  By itself, will attempt to
     generate field, read fields (method defined by derived classes)
@@ -59,13 +61,14 @@ class YTDataContainer:
 
     _chunk_info = None
     _num_ghost_zones = 0
-    _con_args = ()
+    _con_args: Tuple[str, ...] = ()
     _skip_add = False
-    _container_fields = ()
-    _tds_attrs = ()
-    _tds_fields = ()
+    _container_fields: Tuple[Union[str, Tuple[str, str]], ...] = ()
+    _tds_attrs: Tuple[str, ...] = ()
+    _tds_fields: Tuple[str, ...] = ()
     _field_cache = None
     _index = None
+    _key_fields: List[str]
 
     def __init__(self, ds, field_parameters):
         """
@@ -428,8 +431,6 @@ def _parameter_iterate(self, seq):
             yield obj
             obj.field_parameters = old_fp
 
-    _key_fields = None
-
     def write_out(self, filename, fields=None, format="%0.16e"):
         """Write out the YTDataContainer object in a text file.
 
@@ -469,9 +470,6 @@ def write_out(self, filename, fields=None, format="%0.16e"):
         if fields is None:
             fields = sorted(self.field_data.keys())
 
-        if self._key_fields is None:
-            raise ValueError
-
         field_order = [("index", k) for k in self._key_fields]
         diff_fields = [field for field in fields if field not in field_order]
         field_order += diff_fields

yt/data_objects/image_array.py

@@ -1,13 +1,13 @@
 import warnings
 
 import numpy as np
+from unyt import unyt_array
 
 from yt.config import ytcfg
-from yt.units.yt_array import YTArray
 from yt.visualization.image_writer import write_bitmap, write_image
 
 
-class ImageArray(YTArray):
+class ImageArray(unyt_array):
     r"""A custom Numpy ndarray used for images.
 
     This differs from ndarray in that you can optionally specify an

yt/data_objects/index_subobjects/grid_patch.py

@@ -282,7 +282,7 @@ def get_vertex_centered_data(
                 "a single field as an argument."
             )
             warnings.warn(message, DeprecationWarning, stacklevel=2)
-            fields = [fields]
+            fields = [fields]  # type: ignore
 
         # Make sure the field list has only unique entries
         fields = list(set(fields))

yt/data_objects/index_subobjects/octree_subset.py

@@ -1,16 +1,17 @@
 from contextlib import contextmanager
 from itertools import product, repeat
+from typing import Tuple
 
 import numpy as np
+from unyt import unyt_array
 
 import yt.geometry.particle_deposit as particle_deposit
 import yt.geometry.particle_smooth as particle_smooth
 from yt.data_objects.selection_objects.data_selection_objects import (
     YTSelectionContainer,
 )
 from yt.geometry.particle_oct_container import ParticleOctreeContainer
-from yt.units.dimensions import length
-from yt.units.yt_array import YTArray
+from yt.units.dimensions import length  # type: ignore
 from yt.utilities.exceptions import (
     YTFieldTypeNotFound,
     YTInvalidPositionArray,
@@ -37,7 +38,7 @@ class OctreeSubset(YTSelectionContainer):
     _num_ghost_zones = 0
     _type_name = "octree_subset"
     _skip_add = True
-    _con_args = ("base_region", "domain", "ds")
+    _con_args: Tuple[str, ...] = ("base_region", "domain", "ds")
     _domain_offset = 0
     _cell_count = -1
     _block_order = "C"
@@ -658,7 +659,7 @@ def __iter__(self):
             yield i, OctreeSubsetBlockSlicePosition(i, self)
 
 
-class YTPositionArray(YTArray):
+class YTPositionArray(unyt_array):
     @property
     def morton(self):
         self.validate()

yt/data_objects/particle_filters.py

@@ -1,12 +1,13 @@
 import copy
 from contextlib import contextmanager
+from typing import Dict
 
 from yt.fields.field_info_container import NullFunc, TranslationFunc
 from yt.funcs import mylog
 from yt.utilities.exceptions import YTIllDefinedFilter
 
 # One to one mapping
-filter_registry = {}
+filter_registry: Dict[str, "ParticleFilter"] = {}
 
 
 class DummyFieldInfo:

yt/data_objects/profiles.py

@@ -5,7 +5,7 @@
 from yt.fields.derived_field import DerivedField
 from yt.frontends.ytdata.utilities import save_as_dataset
 from yt.funcs import get_output_filename, is_sequence, iter_fields, mylog
-from yt.units.unit_object import Unit
+from yt.units.unit_object import Unit  # type: ignore
 from yt.units.yt_array import YTQuantity, array_like_field
 from yt.utilities.exceptions import (
     YTIllDefinedBounds,

yt/data_objects/selection_objects/data_selection_objects.py

@@ -1,3 +1,4 @@
+import abc
 import itertools
 import uuid
 from collections import defaultdict
@@ -31,13 +32,13 @@
 )
 
 
-class YTSelectionContainer(YTDataContainer, ParallelAnalysisInterface):
+class YTSelectionContainer(YTDataContainer, ParallelAnalysisInterface, abc.ABC):
     _locked = False
     _sort_by = None
     _selector = None
     _current_chunk = None
     _data_source = None
-    _dimensionality = None
+    _dimensionality: int
     _max_level = None
     _min_level = None
     _derived_quantity_chunking = "io"

yt/data_objects/static_output.py

@@ -3,12 +3,14 @@
 import itertools
 import os
 import pickle
+import sys
 import time
 import warnings
 import weakref
 from collections import defaultdict
 from importlib.util import find_spec
 from stat import ST_CTIME
+from typing import Optional, Tuple, Type, Union
 
 import numpy as np
 from unyt.exceptions import UnitConversionError, UnitParseError
@@ -32,11 +34,15 @@
     SpectralCubeCoordinateHandler,
     SphericalCoordinateHandler,
 )
+from yt.geometry.geometry_handler import Index
 from yt.units import UnitContainer, _wrap_display_ytarray, dimensions
-from yt.units.dimensions import current_mks
-from yt.units.unit_object import Unit, define_unit
-from yt.units.unit_registry import UnitRegistry
-from yt.units.unit_systems import create_code_unit_system, unit_system_registry
+from yt.units.dimensions import current_mks  # type: ignore
+from yt.units.unit_object import Unit, define_unit  # type: ignore
+from yt.units.unit_registry import UnitRegistry  # type: ignore
+from yt.units.unit_systems import (  # type: ignore
+    create_code_unit_system,
+    unit_system_registry,
+)
 from yt.units.yt_array import YTArray, YTQuantity
 from yt.utilities.cosmology import Cosmology
 from yt.utilities.exceptions import (
@@ -51,12 +57,18 @@
 from yt.utilities.parallel_tools.parallel_analysis_interface import parallel_root_only
 from yt.utilities.parameter_file_storage import NoParameterShelf, ParameterFileStore
 
+if sys.version_info >= (3, 9):
+    from collections.abc import MutableMapping
+else:
+    from typing import MutableMapping
 # We want to support the movie format in the future.
 # When such a thing comes to pass, I'll move all the stuff that is constant up
 # to here, and then have it instantiate EnzoDatasets as appropriate.
 
 
-_cached_datasets = weakref.WeakValueDictionary()
+_cached_datasets: MutableMapping[
+    Union[int, str], "Dataset"
+] = weakref.WeakValueDictionary()
 _ds_store = ParameterFileStore()
 
 
@@ -114,20 +126,20 @@ class Dataset(abc.ABC):
 
     default_fluid_type = "gas"
     default_field = ("gas", "density")
-    fluid_types = ("gas", "deposit", "index")
-    particle_types = ("io",)  # By default we have an 'all'
-    particle_types_raw = ("io",)
+    fluid_types: Tuple[str, ...] = ("gas", "deposit", "index")
+    particle_types: Optional[Tuple[str, ...]] = ("io",)  # By default we have an 'all'
+    particle_types_raw: Optional[Tuple[str, ...]] = ("io",)
     geometry = "cartesian"
     coordinates = None
     storage_filename = None
     particle_unions = None
     known_filters = None
-    _index_class = None
+    _index_class: Type[Index]
     field_units = None
     derived_field_list = requires_index("derived_field_list")
     fields = requires_index("fields")
     _instantiated = False
-    _unique_identifier = None
+    _unique_identifier: Optional[Union[str, int]] = None
     _particle_type_counts = None
     _proj_type = "quad_proj"
     _ionization_label_format = "roman_numeral"