-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
134 add inputs transformer tools for earthkit subpackages (#148)
* added module_inputs_wrapper and removed unused duplicate script * possibility to escape wrapping at call
- Loading branch information
Showing
6 changed files
with
402 additions
and
179 deletions.
There are no files selected for viewing
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
Oops, something went wrong.