Merge pull request #81 from StagPython/type_annotations

Improve type soundness

This PR adds static type checking with `mypy`. It also comes with a few API changes to make it simpler/more correct wrt to types:

* use of `MappingProxyType` instead of `OrderedDict`
* new `datatypes` module to hold `Var*` as well as `Field`, `Rprof`, and `Tseries` classes
* those classes are now `NamedTuple`s instead of bare `namedtuple`s
* `processing` functions now return `Field`, `Rprof` and `Tseries` instances directly
* new `NoGeomError`, `NoRefstateError` and `NoTimeError` exceptions to signal missing data instead of a plain `None`
* new parsing function `field_header` and `field_istep`. The `fields` function always return the full data instead of having a return type depending on logical arguments
* successive calls to `_StepsView.filter` on the same instance now compose instead of overriding

Author: amorison

.lgtm.yml

@@ -0,0 +1,2 @@
+queries:
+  - exclude: py/import-and-import-from

docs/index.rst

@@ -42,6 +42,7 @@ Welcome to StagPy's documentation!
    sources/apiref/args
    sources/apiref/commands
    sources/apiref/config
+   sources/apiref/datatypes
    sources/apiref/error
    sources/apiref/field
    sources/apiref/parfile

docs/sources/apiref/datatypes.rst

@@ -0,0 +1,11 @@
+datatypes
+=========
+
+.. automodule:: stagpy.datatypes
+
+   .. autoclass:: Varf
+   .. autoclass:: Field
+   .. autoclass:: Varr
+   .. autoclass:: Rprof
+   .. autoclass:: Vart
+   .. autoclass:: Tseries

docs/sources/apiref/phyvars.rst

@@ -3,16 +3,6 @@ phyvars
 
 .. automodule:: stagpy.phyvars
 
-   .. class:: Varf
-
-      :class:`collections.namedtuple` whose instances hold metadata of
-      scalar fields. It defines the following fields:
-
-      - **description** (*str* or *func*): short description of the variable if
-        it is output by StagYY, function to compute it otherwise.
-      - **dim** (*str*): dimension used to
-        :func:`~stagpy.stagyydata.StagyyData.scale` to dimensional values.
-
    .. data:: FIELD
       :annotation: = {fieldvar: Varf()}
 
@@ -31,18 +21,6 @@ phyvars
       Dictionary of surface scalar fields output by StagYY. Keys are the
       variable names, values are :class:`Varf` instances.
 
-   .. class:: Varr
-
-      :class:`collections.namedtuple` whose instances hold metadata of
-      radial profiles. It defines the following fields:
-
-      - **description** (*str* or *func*): short description of the variable if
-        it is output by StagYY, function to compute it otherwise.
-      - **kind** (*str*): shorter description to group similar variables under
-        the same label.
-      - **dim** (*str*): dimension used to
-        :func:`~stagpy.stagyydata.StagyyData.scale` to dimensional values.
-
    .. data:: RPROF
       :annotation: = {rprofvar: Varr()}
 
@@ -61,18 +39,6 @@ phyvars
       Dictionary of radial profiles of the reference state. Keys are the
       variable names, values are :class:`Varr` instances.
 
-   .. class:: Vart
-
-      :class:`collections.namedtuple` whose instances hold metadata of
-      time series. It defines the following fields:
-
-      - **description** (*str* or *func*): short description of the variable if
-        it is output by StagYY, function to compute it otherwise.
-      - **kind** (*str*): shorter description to group similar variables under
-        the same label.
-      - **dim** (*str*): dimension used to
-        :func:`~stagpy.stagyydata.StagyyData.scale` to dimensional values.
-
    .. data:: TIME
       :annotation: = {timevar: Vart()}
 

docs/sources/apiref/step.rst

@@ -4,4 +4,5 @@ _step
 .. automodule:: stagpy._step
    :members:
    :private-members:
-   :exclude-members: _init_shape, _get_raw_data, _scale_radius_mo
+   :exclude-members: _init_shape, _get_raw_data, _scale_radius_mo,
+                     _present_fields

pyproject.toml

@@ -3,3 +3,18 @@ requires = ["setuptools>=45", "setuptools_scm>=6.2", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools_scm]
+
+[tool.mypy]
+check_untyped_defs = true
+
+[[tool.mypy.overrides]]
+module = [
+    "setuptools_scm",
+    "f90nml.*",
+    "h5py.*",
+    "matplotlib.*",
+    "mpl_toolkits.*",
+    "pandas.*",
+    "scipy.*",
+]
+ignore_missing_imports = true

stagpy/__init__.py

@@ -16,7 +16,7 @@
 and uppercase versions of those.
 """
 
-import importlib
+from __future__ import annotations
 import os
 import pathlib
 import shutil
@@ -30,7 +30,7 @@
 from . import config
 
 
-def _env(var):
+def _env(var: str) -> bool:
     """Return whether var is set to True."""
     val = os.getenv(var, default='').lower()
     return val in ('true', 't', 'yes', 'y', 'on', '1')
@@ -69,7 +69,7 @@ def _check_config():
 
 def load_mplstyle():
     """Try to load conf.plot.mplstyle matplotlib style."""
-    plt = importlib.import_module('matplotlib.pyplot')
+    import matplotlib.pyplot as plt
     if conf.plot.mplstyle:
         for style in conf.plot.mplstyle.split():
             found = False

stagpy/__main__.py

@@ -1,6 +1,5 @@
 """The stagpy module is callable."""
 
-import importlib
 import signal
 import sys
 import warnings
@@ -13,8 +12,7 @@ def main():
     if not DEBUG:
         signal.signal(signal.SIGINT, sigint_handler)
         warnings.simplefilter('ignore')
-    args = importlib.import_module('stagpy.args')
-    error = importlib.import_module('stagpy.error')
+    from . import args, error
     try:
         args.parse_args()()
     except error.StagpyError as err:

stagpy/_helpers.py

@@ -1,25 +1,31 @@
 """Various helper functions and classes."""
 
+from __future__ import annotations
 from inspect import getdoc
+from typing import TYPE_CHECKING, Generic, TypeVar
 
 import matplotlib.pyplot as plt
 
 from . import conf
 
+if TYPE_CHECKING:
+    from typing import Optional, Any, List, Callable
+    from matplotlib.figure import Figure
+    from numpy import ndarray
 
-def out_name(stem, timestep=None):
+
+def out_name(stem: str, timestep: Optional[int] = None) -> str:
     """Return StagPy out file name.
 
     Args:
-        stem (str): short description of file content.
-        timestep (int): timestep if relevant.
+        stem: short description of file content.
+        timestep: timestep if relevant.
 
     Returns:
-        str: the output file name.
+        the output file name.
 
     Other Parameters:
-        conf.core.outname (str): the generic name stem, defaults to
-            ``'stagpy'``.
+        conf.core.outname: the generic name stem, defaults to ``'stagpy'``.
     """
     if conf.core.shortname:
         return conf.core.outname
@@ -28,32 +34,33 @@ def out_name(stem, timestep=None):
     return conf.core.outname + '_' + stem
 
 
-def scilabel(value, precision=2):
+def scilabel(value: float, precision: int = 2) -> str:
     """Build scientific notation of some value.
 
     This is dedicated to use in labels displaying scientific values.
 
     Args:
-        value (float): numeric value to format.
-        precision (int): number of decimal digits.
+        value: numeric value to format.
+        precision: number of decimal digits.
 
     Returns:
-        str: the scientific notation the specified value.
+        the scientific notation of the specified value.
     """
-    man, exp = f'{value:.{precision}e}'.split('e')
-    exp = int(exp)
+    man, exps = f'{value:.{precision}e}'.split('e')
+    exp = int(exps)
     return fr'{man}\times 10^{{{exp}}}'
 
 
-def saveplot(fig, *name_args, close=True, **name_kwargs):
+def saveplot(fig: Figure, *name_args: Any, close: bool = True,
+             **name_kwargs: Any):
     """Save matplotlib figure.
 
     You need to provide :data:`stem` as a positional or keyword argument (see
     :func:`out_name`).
 
     Args:
-        fig (:class:`matplotlib.figure.Figure`): matplotlib figure.
-        close (bool): whether to close the figure.
+        fig: the :class:`matplotlib.figure.Figure` to save.
+        close: whether to close the figure.
         name_args: positional arguments passed on to :func:`out_name`.
         name_kwargs: keyword arguments passed on to :func:`out_name`.
     """
@@ -64,7 +71,7 @@ def saveplot(fig, *name_args, close=True, **name_kwargs):
         plt.close(fig)
 
 
-def baredoc(obj):
+def baredoc(obj: object) -> str:
     """Return the first line of the docstring of an object.
 
     Trailing periods and spaces as well as leading spaces are removed from the
@@ -82,12 +89,12 @@ def baredoc(obj):
     return doc.rstrip(' .').lstrip()
 
 
-def list_of_vars(arg_plot):
+def list_of_vars(arg_plot: str) -> List[List[List[str]]]:
     """Construct list of variables per plot.
 
     Args:
-        arg_plot (str): string with variable names separated with
-            ``-`` (figures), ``.`` (subplots) and ``,`` (same subplot).
+        arg_plot: variable names separated with ``-`` (figures),
+            ``.`` (subplots) and ``,`` (same subplot).
     Returns:
         three nested lists of str
 
@@ -102,13 +109,13 @@ def list_of_vars(arg_plot):
     return [lov for lov in lovs if lov]
 
 
-def find_in_sorted_arr(value, array, after=False):
+def find_in_sorted_arr(value: Any, array: ndarray, after=False) -> int:
     """Return position of element in a sorted array.
 
     Returns:
-        int: the maximum position i such as array[i] <= value.  If after is
-            True, it returns the min i such as value <= array[i] (or 0 if such
-            an indices does not exist).
+        the maximum position i such as array[i] <= value.  If after is True, it
+        returns the min i such as value <= array[i] (or 0 if such an index does
+        not exist).
     """
     ielt = array.searchsorted(value)
     if ielt == array.size:
@@ -118,7 +125,11 @@ def find_in_sorted_arr(value, array, after=False):
     return ielt
 
 
-class CachedReadOnlyProperty:
+T = TypeVar('T')
+V = TypeVar('V')
+
+
+class CachedReadOnlyProperty(Generic[T, V]):
     """Descriptor implementation of read-only cached properties.
 
     Properties are cached as ``_cropped_{name}`` instance attribute.
@@ -133,13 +144,13 @@ class CachedReadOnlyProperty:
     property is read-only instead of being writeable.
     """
 
-    def __init__(self, thunk):
+    def __init__(self, thunk: Callable[[T], V]):
         self._thunk = thunk
         self._name = thunk.__name__
         self._cache_name = f'_cropped_{self._name}'
         self.__doc__ = thunk.__doc__
 
-    def __get__(self, instance, _):
+    def __get__(self, instance: T, _) -> V:
         try:
             return getattr(instance, self._cache_name)
         except AttributeError:
@@ -148,6 +159,6 @@ def __get__(self, instance, _):
         setattr(instance, self._cache_name, cached_value)
         return cached_value
 
-    def __set__(self, instance, _):
+    def __set__(self, instance: T, _):
         raise AttributeError(
             f'Cannot set {self._name} property of {instance!r}')