processing.stream_function returns a Field

It was previously returning an array. This change is to improve the
type soundness of the public API.

To avoid cyclic dependencies, a new `datatypes` is created to hold definitions
of several types that are ubiquitous in the code base.

Author: amorison

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,7 @@
+datatypes
+=========
+
+.. automodule:: stagpy.datatypes
+
+   .. autoclass:: Varf
+   .. autoclass:: Field

docs/sources/apiref/phyvars.rst

@@ -3,8 +3,6 @@ phyvars
 
 .. automodule:: stagpy.phyvars
 
-   .. autoclass:: Varf
-
    .. data:: FIELD
       :annotation: = {fieldvar: Varf()}
 

docs/sources/apiref/step.rst

@@ -5,7 +5,6 @@ _step
    :members:
    :private-members:
    :exclude-members: _init_shape, _get_raw_data, _scale_radius_mo,
-                     Field, Rprof, _present_fields
+                     Rprof, _present_fields
 
-   .. autoclass:: Field
    .. autoclass:: Rprof

stagpy/_step.py

@@ -15,12 +15,15 @@
 
 from . import error, phyvars, stagyyparsers, _helpers
 from ._helpers import CachedReadOnlyProperty as crop
+from .datatypes import Field
 
 if TYPE_CHECKING:
-    from typing import Dict, Any, Mapping, List, Iterator, Tuple, Optional
+    from typing import (Dict, Any, Mapping, List, Iterator, Tuple, Optional,
+                        Callable)
     from numpy import ndarray, signedinteger
     from pandas import DataFrame, Series
-    from .phyvars import Varf, Varr
+    from .datatypes import Varf
+    from .phyvars import Varr
     from .stagyydata import StagyyData
 
 
@@ -226,34 +229,22 @@ def at_r(self, rval: float) -> signedinteger:
         return np.argmin(np.abs(self.r_centers - rval))
 
 
-class Field(NamedTuple):
-    """Scalar field and associated metadata.
-
-    Attributes:
-        values: the field itself.
-        meta: the metadata of the field.
-    """
-
-    values: ndarray
-    meta: Varf
-
-
 class _Fields(abc.Mapping):
     """Fields data structure.
 
     The :attr:`Step.fields` attribute is an instance of this class.
 
     :class:`_Fields` inherits from :class:`collections.abc.Mapping`. Keys are
     fields names defined in :data:`stagpy.phyvars.[S]FIELD[_EXTRA]`.  Each item
-    is a name tuple ('values', 'meta'), respectively the field itself, and a
-    :class:`stagpy.phyvars.Varf` instance with relevant metadata.
+    is a :class:`stagpy.datatypes.Field` instance.
 
     Attributes:
         step: the step object owning the :class:`_Fields` instance.
     """
 
     def __init__(self, step: Step, variables: Mapping[str, Varf],
-                 extravars: Mapping[str, Varf], files: Mapping[str, List[str]],
+                 extravars: Mapping[str, Callable[[Step], Field]],
+                 files: Mapping[str, List[str]],
                  filesh5: Mapping[str, List[str]]):
         self.step = step
         self._vars = variables
@@ -269,10 +260,7 @@ def __getitem__(self, name: str) -> Field:
         if name in self._vars:
             fld_names, parsed_data = self._get_raw_data(name)
         elif name in self._extra:
-            meta = self._extra[name]
-            field = meta.description(self.step)
-            meta = phyvars.Varf(_helpers.baredoc(meta.description), meta.dim)
-            self._data[name] = Field(field, meta)
+            self._data[name] = self._extra[name](self.step)
             return self._data[name]
         else:
             raise error.UnknownFieldVarError(name)

stagpy/commands.py

@@ -110,8 +110,7 @@ def _pretty_print(key_val, sep=': ', min_col_width=39, text_width=None):
 def _layout(dict_vars, dict_vars_extra):
     """Print nicely [(var, description)] from phyvars."""
     desc = [(v, m.description) for v, m in dict_vars.items()]
-    desc.extend((v, baredoc(m.description))
-                for v, m in dict_vars_extra.items())
+    desc.extend((v, baredoc(m)) for v, m in dict_vars_extra.items())
     _pretty_print(desc, min_col_width=26)
 
 

stagpy/datatypes.py

@@ -0,0 +1,32 @@
+"""Types describing StagYY output data."""
+
+from __future__ import annotations
+from typing import NamedTuple, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from numpy import ndarray
+
+
+class Varf(NamedTuple):
+    """Metadata of scalar field.
+
+    Attributes:
+        description: short description of the variable.
+        dim: dimension used to :func:`~stagpy.stagyydata.StagyyData.scale` to
+            dimensional values.
+    """
+
+    description: str
+    dim: str
+
+
+class Field(NamedTuple):
+    """Scalar field and associated metadata.
+
+    Attributes:
+        values: the field itself.
+        meta: the metadata of the field.
+    """
+
+    values: ndarray
+    meta: Varf

stagpy/phyvars.py

@@ -11,6 +11,7 @@
 from typing import NamedTuple, TYPE_CHECKING
 
 from . import processing
+from .datatypes import Varf
 
 if TYPE_CHECKING:
     from typing import Union, Callable, Tuple
@@ -19,20 +20,6 @@
     from .stagyydata import StagyyData
 
 
-class Varf(NamedTuple):
-    """Metadata of scalar fields.
-
-    Attributes:
-        description: short description of the variable if it is output by
-            StagYY, function to compute it otherwise.
-        dim: dimension used to :func:`~stagpy.stagyydata.StagyyData.scale` to
-            dimensional values.
-    """
-
-    description: Union[str, Callable[[Step], ndarray]]
-    dim: str
-
-
 FIELD = MappingProxyType({
     'T': Varf('Temperature', 'K'),
     'v1': Varf('x Velocity', 'm/s'),
@@ -66,7 +53,7 @@ class Varf(NamedTuple):
 })
 
 FIELD_EXTRA = MappingProxyType({
-    'stream': Varf(processing.stream_function, 'm2/s'),
+    'stream': processing.stream_function,
 })
 
 FIELD_FILES = MappingProxyType({

stagpy/processing.py

@@ -12,6 +12,7 @@
 from scipy import integrate
 
 from .error import NotAvailableError
+from .datatypes import Field, Varf
 
 if typing.TYPE_CHECKING:
     from typing import Tuple
@@ -277,13 +278,13 @@ def c_overturned(step: Step) -> Tuple[ndarray, ndarray]:
     return cinit, radf
 
 
-def stream_function(step: Step) -> ndarray:
-    """Stream function.
+def stream_function(step: Step) -> Field:
+    """Compute the stream function in 2D geometry.
 
     Args:
         step: a :class:`~stagpy._step.Step` of a StagyyData instance.
     Returns:
-        the stream function field, with four dimensions (x, y, z and block).
+        the stream function field.
     """
     if step.geom.twod_yz:
         x_coord = step.geom.y_walls
@@ -327,4 +328,4 @@ def stream_function(step: Step) -> ndarray:
     if step.geom.twod_xz:
         psi = - psi
     psi = np.reshape(psi, shape)
-    return psi
+    return Field(psi, Varf("Stream function", 'm2/s'))

tests/test_phyvars.py

@@ -3,7 +3,7 @@
 
 
 def test_dim():
-    allvars = chain(phyvars.FIELD.values(), phyvars.FIELD_EXTRA.values(),
+    allvars = chain(phyvars.FIELD.values(),
                     phyvars.RPROF.values(), phyvars.RPROF_EXTRA.values(),
                     phyvars.TIME.values(), phyvars.TIME_EXTRA.values())
     for var in allvars: