CLN: add typing to _gridprop_value_init

Includes some mild refactoring.

src/xtgeo/grid3d/_gridprop_value_init.py

@@ -1,101 +1,147 @@
 """GridProperty (not GridProperies) some etc functions"""
+from __future__ import annotations
 
-# from ._gridprop_import_eclrun import import_eclbinary
-#  --> INIT, UNRST
-#
-# from ._gridprop_import_grdecl import import_grdecl_prop, import_bgrdecl_prop
-#  --> ASCII and BINARY GRDECL format
-# from ._gridprop_import_roff import import_roff
-#  --> BINARY ROFF format
-
+import numbers
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import numpy as np
 
 import xtgeo
+from xtgeo.common import XTGeoDialog
 
-xtg = xtgeo.common.XTGeoDialog()
-
+xtg = XTGeoDialog()
 logger = xtg.functionlogger(__name__)
 
+if TYPE_CHECKING:
+    from xtgeo.grid3d import Grid, GridProperty
 
-def gridproperty_non_dummy_values(gridlike, dimensions, values, isdiscrete):
-    """Gives the initial values array of an gridprop.
 
-    Param:
-      gridlike: Either Grid or GridProperty, giving the mask to replicate.
-      dimensions: The dimensions of the gridprop
-      values: The values parameter given to init.
-      isdiscrete: The discrete parameter given to init.
+def gridproperty_non_dummy_values(
+    gridlike: Optional[Union[Grid, GridProperty]],
+    dimensions: Tuple[int, int, int],
+    values: Optional[Union[np.ndarray, float, int]],
+    isdiscrete: bool,
+) -> np.ma.MaskedArray:
+    """
+    Gives the initial values array of an gridprop.
+
+    Parameters:
+        gridlike:
+            Either Grid or GridProperty, giving the mask to replicate.
+        dimensions:
+            The (ncol, nrow, nlay) dimensions of the grid property.
+        values:
+            The values parameter given to init.
+        isdiscrete:
+            The discrete parameter given to init.
 
     Returns:
-      The array to be set to GridProp._values.
+        The array to be set to GridProperty._values.
+
     """
     if values is None:
-        values = initial_gridprop_values_zero(dimensions, isdiscrete)
-    elif np.isscalar(values):
-        values = initial_gridprop_values_from_scalar(dimensions, values, isdiscrete)
+        _values = initial_gridprop_values_zero(dimensions, isdiscrete)
+    elif isinstance(values, numbers.Number):
+        _values = initial_gridprop_values_from_scalar(dimensions, values, isdiscrete)
     elif isinstance(values, np.ndarray):
-        values = initial_gridprop_values_from_array(dimensions, values, isdiscrete)
+        _values = initial_gridprop_values_from_array(dimensions, values, isdiscrete)
 
-    if gridlike is not None:
+    if gridlike:
         if isinstance(gridlike, xtgeo.grid3d.Grid):
             act = gridlike.get_actnum(asmasked=True)
-            values = np.ma.array(values, mask=np.ma.getmaskarray(act.values))
+            _values = np.ma.array(_values, mask=np.ma.getmaskarray(act.values))
         else:
-            values = np.ma.array(values, mask=np.ma.getmaskarray(gridlike.values))
+            assert isinstance(gridlike, xtgeo.grid3d.GridProperty)
+            _values = np.ma.array(_values, mask=np.ma.getmaskarray(gridlike.values))
+
+    return _values
 
-    return values
 
+def gridproperty_dummy_values(isdiscrete: bool) -> np.ma.MaskedArray:
+    """
+    Given no parameters to init, these dummy values should be set for backwards
+    compatability.
+
+    Parameters:
+        isdiscrete:
+            If the grid property values are discrete.
+
+    Returns:
+        The array to be set to GridProperty._values.
 
-def gridproperty_dummy_values(isdiscrete):
-    """Given no parameters to init, these dummy values should be set for backwards
-    compatability."""
-    if isdiscrete:
-        values = np.ma.MaskedArray(np.full((4, 3, 5), 99), dtype=np.int32)
-    else:
-        values = np.ma.MaskedArray(np.full((4, 3, 5), 99.0))
+    """
+    values: np.ma.MaskedArray = np.ma.MaskedArray(
+        np.full((4, 3, 5), 99.0), dtype=np.int32 if isdiscrete else np.float64
+    )
     values[0:4, 0, 0:2] = np.ma.masked
     return values
 
 
-def initial_gridprop_values_zero(dimensions, isdiscrete):
-    """Initial values for an GridProperty with zeros.
+def initial_gridprop_values_zero(
+    dimensions: Tuple[int, int, int], isdiscrete: bool
+) -> np.ma.MaskedArray:
+    """
+    Initial values for an GridProperty with zeros.
 
     Given that the user supplies at least some parameters, but not a values array,
     values should be initialized to zero.
-    Param:
-      dimensions: The dimensions of the gridproperty.
 
+    Parameters:
+        dimensions:
+            The (ncol, nrow, nlay) dimensions of the grid property.
 
     Returns:
-        zero initialized values array
+        Zero initialized values array.
+
     """
-    if isdiscrete:
-        return np.ma.zeros(dimensions, dtype=np.int32)
-    return np.ma.zeros(dimensions)
+    return np.ma.zeros(dimensions, dtype=np.int32 if isdiscrete else np.float64)
 
 
-def initial_gridprop_values_from_scalar(dimensions, value, isdiscrete):
-    """Initial gridproperties values from scalar.
+def initial_gridprop_values_from_scalar(
+    dimensions: Tuple[int, int, int], value: Union[float, int], isdiscrete: bool
+) -> np.ma.MaskedArray:
+    """
+    Initial grid property values from scalar.
 
     Given scalar values, the gridproperties value array should be
     filled with that value, with possible conversion depending
     on the isdiscrete parameter.
 
+    Parameters:
+        dimensions:
+            The (ncol, nrow, nlay) dimensions of the grid property.
+        value:
+            The scalar value to initialize with.
+        isdiscrete:
+            If the values are discrete.
+
+    Returns:
+        Filled array with given scalar value.
+
+    """
+    if not isinstance(value, numbers.Number):
+        raise ValueError("Scalar input values of invalid type")
+    return np.ma.zeros(dimensions, dtype=np.int32 if isdiscrete else np.float64) + value
+
+
+def initial_gridprop_values_from_array(
+    dimensions: Tuple[int, int, int], values: np.ndarray, isdiscrete: bool
+) -> np.ma.MaskedArray:
+    """
+    Initial GridProperty values from numpy array.
+
+    Parameters:
+        dimensions:
+            The (ncol, nrow, nlay) dimensions of the grid property.
+        value:
+            The numpy array to initialize with.
+        isdiscrete:
+            If the values are discrete.
+
     Returns:
-        filled array with given scalar value.
+        GridProperty with values initialized from a numpy array.
+
     """
-    if isinstance(value, (float, int)):
-        dtype = np.float64
-        if isdiscrete:
-            dtype = np.int32
-        return np.ma.zeros(dimensions, dtype=dtype) + value
-    raise ValueError("Scalar input values of invalid type")
-
-
-def initial_gridprop_values_from_array(dimensions, values, isdiscrete):
-    """Initial gridproperties values from numpy array"""
-    values = np.ma.MaskedArray(values.reshape(dimensions))
-    if isdiscrete:
-        return values.astype(np.int32)
-    return values.astype(np.float64)
+    return np.ma.MaskedArray(
+        values.reshape(dimensions), dtype=np.int32 if isdiscrete else np.float64
+    )