134 add inputs transformer tools for earthkit subpackages (#148)

* added module_inputs_wrapper and removed unused duplicate script
 
* possibility to escape wrapping at call

earthkit/data/utils/module_inputs_wrapper.py

@@ -0,0 +1,197 @@
+# (C) Copyright 2020 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+"""
+Module containing methods to transform the inputs of functions based on the function type setting,
+common signitures or mapping defined at call time
+"""
+import inspect
+import types
+import typing as T
+from ast import Module
+from functools import wraps
+
+from earthkit.data import transform
+from earthkit.data.wrappers import Wrapper
+
+try:
+    UNION_TYPES = [T.Union, types.UnionType]
+except AttributeError:
+    # types.UnionType is not in versions of python<3.9
+    UNION_TYPES = [
+        T.Union,
+    ]
+
+EMPTY_TYPES = [inspect._empty]
+
+
+def _ensure_iterable(input_item):
+    """Ensure that an item is iterable"""
+    if not isinstance(input_item, (tuple, list, dict)):
+        return [input_item]
+    return input_item
+
+
+def _ensure_tuple(input_item):
+    """Ensure that an item is a tuple"""
+    if not isinstance(input_item, tuple):
+        return tuple(_ensure_iterable(input_item))
+    return input_item
+
+
+def transform_function_inputs(
+    function: T.Callable,
+    kwarg_types: T.Dict[str, T.Any] = {},
+    convert_types: T.Union[T.Tuple[T.Any], T.Dict[str, T.Tuple[T.Any]]] = (),
+) -> T.Callable:
+    """
+    Transform the inputs to a function to match the requirements.
+    earthkit.data handles the input arg/kwarg format.
+
+    Parameters
+    ----------
+    function : Callable
+        Method to be wrapped
+    kwarg_types : Dict[str: type]
+        Mapping of accepted object types for each arg/kwarg
+    convert_types : Tuple[type]
+        List of data-types to try to convert, this can be useful when the function is versitile and can
+        accept a large number of data-types, hence only a small number of types should be converted.
+
+    Returns
+    -------
+    [type]
+        [description]
+    """
+
+    def _wrapper(kwarg_types, convert_types, *args, **kwargs):
+        kwarg_types = {**kwarg_types}
+        signature = inspect.signature(function)
+        mapping = signature_mapping(signature, kwarg_types)
+
+        # Add args to kwargs for ease of looping:
+        arg_names = []
+        for arg, name in zip(args, signature.parameters):
+            arg_names.append(name)
+            kwargs[name] = arg
+
+        convert_kwargs = [k for k in kwargs if k in mapping]
+        # Only convert some data-types, this can be used to prevent conversion for for functions which
+        #  accept a long-list of formats, e.g. numpy methods can accept xarray, pandas and more
+
+        # Filter for convert_types
+        if convert_types:
+            # Ensure convert_types is a dictionary
+            if not isinstance(convert_types, dict):
+                convert_types = {key: convert_types for key in convert_kwargs}
+
+            convert_kwargs = [
+                k
+                for k in convert_kwargs
+                if isinstance(kwargs[k], _ensure_tuple(convert_types.get(k, ())))
+            ]
+
+        # transform args/kwargs
+        for key in convert_kwargs:
+            value = kwargs[key]
+            kwarg_types = _ensure_iterable(mapping[key])
+            # Transform value if necessary
+            if type(value) not in kwarg_types:
+                for kwarg_type in kwarg_types:
+                    try:
+                        kwargs[key] = transform(value, kwarg_type)
+                    except Exception:
+                        # Transform was not possible, move to next kwarg type.
+                        # If no transform is possible, format is unchanged and we rely on function to raise
+                        # an Error.
+                        continue
+                    break
+
+        # Anything that is still a Wrapper object, expand to native data format:
+        for k, v in [(_k, _v) for _k, _v in kwargs.items() if isinstance(_v, Wrapper)]:
+            try:
+                kwargs[k] = v.data
+            except Exception:
+                pass
+
+        # Extract args from kwargs:
+        args = [kwargs.pop(name) for name in arg_names]
+        return function(*args, **kwargs)
+
+    @wraps(function)
+    def wrapper(*args, _auto_inputs_transform=True, **kwargs):
+        if not _auto_inputs_transform:  # Possibility to escape wrapping at call level
+            return function(*args, **kwargs)
+        return _wrapper(kwarg_types, convert_types, *args, **kwargs)
+
+    return wrapper
+
+
+def signature_mapping(signature, kwarg_types):
+    """
+    Map args and kwargs to object types, using hierarchical selection method:
+    1. Explicitly defined type
+    2. Based on Type setting in function
+    3. Do nothing
+    """
+    mapping = {}
+    for key, parameter in signature.parameters.items():
+        if key in kwarg_types:
+            # 1. Use explicitly defined type
+            kwarg_type = kwarg_types.get(key)
+        elif parameter.annotation not in EMPTY_TYPES:
+            # 2. Use type setting from function
+            kwarg_type = parameter.annotation
+            if T.get_origin(kwarg_type) in UNION_TYPES:
+                # Need to expand union_types to list
+                kwarg_type = T.get_args(kwarg_type)
+        else:
+            # 3. Do nothing, cannot assign None, as None is a valid type
+            continue
+        mapping[key] = kwarg_type
+    return mapping
+
+
+def transform_module_inputs(in_module, **kwargs):
+    """
+    Transform the inputs to all functions in a module.
+
+    Parameters
+    ----------
+    in_module : Module
+        Module containing funcitons which are to be wrapped with transform_function_inputs
+    kwarg_types : Dict[str: type]
+        Mapping of accepted object types for each arg/kwarg
+    convert_types : Tuple[type]
+        List of data-types to try to convert, this can be useful when the function is versitile and can
+        accept a large number of data-types, hence only a small number of types should be converted.
+    kwargs: Any
+        Any other kwargs accepted by transform_function_inputs
+
+    Returns
+    -------
+    Module
+        Version of in_module where all functions are wrapped with transform_modul_inputs
+    """
+    # wrapped_module must be different to original to prevent overriding cached module
+    wrapped_module = Module
+    for name in dir(in_module):
+        func = getattr(in_module, name)
+        # Wrap any functions that are not hidden
+        if not name.startswith("_") and isinstance(func, types.FunctionType):
+            setattr(
+                wrapped_module,
+                name,
+                transform_function_inputs(func, **kwargs),
+            )
+        else:
+            # If not a func, we just copy
+            setattr(wrapped_module, name, func)
+
+    return wrapped_module

earthkit/data/wrappers/ndarray.py

@@ -28,6 +28,18 @@ def to_numpy(self):
         """
         return self.data
 
+    def to_xarray(self, **kwargs):
+        """
+        Return an xarray.DataArray representation of the data.
+
+        Returns
+        -------
+        xarray.DataArray
+        """
+        import xarray as xr
+
+        return xr.DataArray(self.data, **kwargs)
+
 
 def wrapper(data, *args, **kwargs):
     import numpy as np