From c9028129954c86458b0d63c6a053da6d3f69855b Mon Sep 17 00:00:00 2001 From: Sandor Kertesz Date: Wed, 11 Mar 2026 10:57:17 +0000 Subject: [PATCH] Enable input modules wrappers tests --- tests/utils/test_module_inputs_wrapper.py | 336 +++++++++++----------- 1 file changed, 168 insertions(+), 168 deletions(-) diff --git a/tests/utils/test_module_inputs_wrapper.py b/tests/utils/test_module_inputs_wrapper.py index 5b8e77e8..00635ded 100644 --- a/tests/utils/test_module_inputs_wrapper.py +++ b/tests/utils/test_module_inputs_wrapper.py @@ -6,173 +6,173 @@ # granted to it by virtue of its status as an intergovernmental organisation # nor does it submit to any jurisdiction. # -import pytest - -pytest.skip(allow_module_level=True) # noqa: E402 - -# import numpy as np -# import pandas as pd -# import xarray as xr - -# from earthkit.data import from_object -# from earthkit.data import from_source -# from earthkit.data.readers import Reader -# from earthkit.data.utils import module_inputs_wrapper -# from earthkit.data.utils.testing import earthkit_test_data_file - -# from . import dummy_module # noqa: F402 -# from .dummy_module import XR_TYPES # noqa: E402 - -# TEST_NP = np.arange(10) -# TEST_NP2 = np.arange(10) - -# TEST_DF = pd.DataFrame({"index": TEST_NP, "data": TEST_NP2}).set_index("index") - -# TEST_DA = xr.DataArray(TEST_NP, name="test") -# TEST_DA2 = xr.DataArray(TEST_NP2, name="test2") -# TEST_DS = TEST_DA.to_dataset() -# TEST_DS["test2"] = TEST_DA2 +import numpy as np +import pandas as pd +import xarray as xr + +from earthkit.data import from_object +from earthkit.data import from_source +from earthkit.data.data import Data +from earthkit.data.readers import Reader +from earthkit.data.utils import module_inputs_wrapper +from earthkit.data.utils.testing import earthkit_test_data_file -# EK_XARRAY_WRAPPER = from_object(TEST_DS) -# EK_NUMPY_WRAPPER = from_object(TEST_NP) - -# WRAPPED_XR_ONES_LIKE = module_inputs_wrapper.transform_function_inputs( -# xr.ones_like, kwarg_types={"other": XR_TYPES} -# ) - -# WRAPPED_XR_ONES_LIKE_TYPE_SETTING = module_inputs_wrapper.transform_function_inputs( -# dummy_module.xarray_ones_like, -# ) +from . import dummy_module # noqa: F402 +from .dummy_module import XR_TYPES # noqa: E402 + +# pytest.skip(allow_module_level=True) # noqa: E402 + + +TEST_NP = np.arange(10) +TEST_NP2 = np.arange(10) + +TEST_DF = pd.DataFrame({"index": TEST_NP, "data": TEST_NP2}).set_index("index") + +TEST_DA = xr.DataArray(TEST_NP, name="test") +TEST_DA2 = xr.DataArray(TEST_NP2, name="test2") +TEST_DS = TEST_DA.to_dataset() +TEST_DS["test2"] = TEST_DA2 -# WRAPPED_NP_MEAN = module_inputs_wrapper.transform_function_inputs( -# np.mean, -# kwarg_types={"a": np.ndarray}, -# convert_types=(Reader), # Only convert Earthkit.data.Reader (np.mean can handle xarray and pandas) -# ) - -# WRAPPED_NP_MEAN_TYPE_SETTING = module_inputs_wrapper.transform_function_inputs( -# dummy_module.numpy_mean, -# ) - -# WRAPPED_DUMMY_MODULE = module_inputs_wrapper.transform_module_inputs( -# dummy_module, -# ) - - -# def test_transform_function_inputs_reader_to_xarray(): -# # Check EK GribReader object -# EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib")) -# ek_reader_result = WRAPPED_XR_ONES_LIKE(EK_GRIB_READER) -# # Will return a DataSet becuase that is first value in kwarg_types -# assert isinstance(ek_reader_result, xr.Dataset) -# assert ek_reader_result.equals(xr.ones_like(EK_GRIB_READER.to_xarray())) - - -# def test_transform_function_inputs_reader_to_xarray_typesetting(): -# # Check EK GribReader object -# EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib")) -# ek_reader_result = WRAPPED_XR_ONES_LIKE_TYPE_SETTING(EK_GRIB_READER) -# # Will return a dataarray because that is first value in type-set Union -# assert isinstance(ek_reader_result, xr.DataArray) -# assert ek_reader_result.equals(xr.ones_like(EK_GRIB_READER.to_xarray()["2t"])) - - -# def test_transform_module_inputs_reader_to_xarray(): -# # Check EK GribReader object -# EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib")) -# ek_reader_result = WRAPPED_DUMMY_MODULE.xarray_ones_like(EK_GRIB_READER) -# # Data array because type-setting of function has dataarray first -# assert isinstance(ek_reader_result, xr.DataArray) -# assert ek_reader_result.equals(xr.ones_like(EK_GRIB_READER.to_xarray())["2t"]) - - -# def test_transform_function_inputs_wrapper_to_xarray(): -# # EK XarrayWrapper object -# ek_wrapper_result = WRAPPED_XR_ONES_LIKE(EK_XARRAY_WRAPPER) -# assert isinstance(ek_wrapper_result, xr.Dataset) -# assert ek_wrapper_result.equals(xr.ones_like(EK_XARRAY_WRAPPER.data)) -# # EK NumpyWrapper object -# ek_wrapper_result = WRAPPED_XR_ONES_LIKE(EK_NUMPY_WRAPPER) -# assert isinstance(ek_wrapper_result, xr.DataArray) -# assert ek_wrapper_result.equals(xr.ones_like(TEST_DA)) - - -# def test_transform_module_inputs_wrapper_to_xarray(): -# # EK XarrayWrapper object -# ek_wrapper_result = WRAPPED_DUMMY_MODULE.xarray_ones_like(EK_XARRAY_WRAPPER) -# # Will return a dataarray because that is first value in type-set Union -# assert isinstance(ek_wrapper_result, xr.DataArray) -# assert ek_wrapper_result.equals(xr.ones_like(EK_XARRAY_WRAPPER.data.test)) -# # EK NumpyWrapper object -# ek_wrapper_result = WRAPPED_DUMMY_MODULE.xarray_ones_like(EK_NUMPY_WRAPPER) -# assert isinstance(ek_wrapper_result, xr.DataArray) -# assert ek_wrapper_result.equals(xr.ones_like(TEST_DA)) - - -# def test_transform_function_inputs_reader_to_numpy(): -# # Test with Earthkit.data GribReader object -# EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib")) -# assert WRAPPED_NP_MEAN(EK_GRIB_READER) == np.mean(EK_GRIB_READER.to_numpy()) -# assert isinstance(WRAPPED_NP_MEAN(EK_GRIB_READER), np.float64) - - -# def test_transform_function_inputs_reader_to_numpy_typesetting(): -# # Test with Earthkit.data GribReader object -# EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib")) -# result = WRAPPED_NP_MEAN_TYPE_SETTING(EK_GRIB_READER) -# assert result == np.mean(EK_GRIB_READER.to_numpy()) -# assert isinstance(result, np.float64) - - -# def test_transform_module_inputs_reader_to_numpy(): -# # Test with Earthkit.data GribReader object -# EK_GRIB_READER = from_source("file", earthkit_test_data_file("test_single.grib")) -# result = WRAPPED_DUMMY_MODULE.numpy_mean(EK_GRIB_READER) -# assert result == np.mean(EK_GRIB_READER.to_numpy()) -# assert isinstance(result, np.float64) - - -# def test_transform_function_inputs_wrapper_to_numpy(): -# # Test with Earthkit.data XarrayWrapper object -# ek_object_result = WRAPPED_NP_MEAN(EK_XARRAY_WRAPPER) -# assert ek_object_result == np.mean(TEST_DS) -# # assert isinstance(ek_object_result, type(EK_XARRAY_WRAPPER.data)) - -# # Test with Earthkit.data NumpyWrapper object -# ek_object_result = WRAPPED_NP_MEAN(EK_NUMPY_WRAPPER) -# assert ek_object_result == np.mean(TEST_NP) -# assert isinstance(ek_object_result, np.float64) - - -# def test_transform_module_inputs_wrapper_to_numpy(): -# # Test with Earthkit.data XarrayWrapper object -# ek_object_result = WRAPPED_DUMMY_MODULE.numpy_mean(EK_XARRAY_WRAPPER) -# assert ek_object_result == np.mean(TEST_DS) -# # TODO: Without definition of convert_types, WRAPPERS are converted to first arguement -# # This could be addressed in future versions -# # assert isinstance(ek_object_result, type(EK_XARRAY_WRAPPER.data)) - -# # Test with Earthkit.data NumpyWrapper object -# ek_object_result = WRAPPED_DUMMY_MODULE.numpy_mean(EK_NUMPY_WRAPPER) -# assert ek_object_result == np.mean(TEST_NP) -# assert isinstance(ek_object_result, np.float64) - - -# def test_transform_function_inputs_xarray_to_numpy(): -# # Test with xarray.DataArray object -# assert WRAPPED_NP_MEAN(TEST_DA) == np.mean(TEST_NP) -# assert WRAPPED_NP_MEAN(TEST_DA) == np.mean(TEST_DA) -# assert isinstance(WRAPPED_NP_MEAN(TEST_DA), xr.DataArray) - - -# def test_transform_function_inputs_pandas_to_numpy(): -# # Test with pandas.DataFrame object, axis=0 is to preserve the pandas.DataFrame, this is pandas syntax -# result = WRAPPED_NP_MEAN(TEST_DF, axis=0) -# assert result.equals(np.mean(TEST_DF, axis=0)) -# assert isinstance(result, pd.core.series.Series) - -# # # Test without axis declartion, pandas>2 returns a scalar mean of the dataframe -# # # Omitted as breaking change in pandas makes tests unstable -# # result = WRAPPED_NP_MEAN(TEST_DF) -# # assert result == np.mean(TEST_DF) -# # assert isinstance(result, np.float64) +EK_XARRAY_WRAPPER = from_object(TEST_DS) +EK_NUMPY_WRAPPER = from_object(TEST_NP) + +WRAPPED_XR_ONES_LIKE = module_inputs_wrapper.transform_function_inputs( + xr.ones_like, kwarg_types={"other": XR_TYPES} +) + +WRAPPED_XR_ONES_LIKE_TYPE_SETTING = module_inputs_wrapper.transform_function_inputs( + dummy_module.xarray_ones_like, +) + +WRAPPED_NP_MEAN = module_inputs_wrapper.transform_function_inputs( + np.mean, + kwarg_types={"a": np.ndarray}, + convert_types=(Data, Reader), # Only convert Earthkit.data.Reader (np.mean can handle xarray and pandas) +) + +WRAPPED_NP_MEAN_TYPE_SETTING = module_inputs_wrapper.transform_function_inputs( + dummy_module.numpy_mean, +) + +WRAPPED_DUMMY_MODULE = module_inputs_wrapper.transform_module_inputs( + dummy_module, +) + + +def test_transform_function_inputs_gribdata_to_xarray(): + # Check EK GribData object + EK_GRIB_DATA = from_source("file", earthkit_test_data_file("test_single.grib")) + ek_reader_result = WRAPPED_XR_ONES_LIKE(EK_GRIB_DATA) + # Will return a DataSet because that is first value in kwarg_types + assert isinstance(ek_reader_result, xr.Dataset) + assert ek_reader_result.equals(xr.ones_like(EK_GRIB_DATA.to_xarray())) + + +def test_transform_function_inputs_gribdata_to_xarray_typesetting(): + # Check EK GribData object + EK_GRIB_DATA = from_source("file", earthkit_test_data_file("test_single.grib")) + ek_reader_result = WRAPPED_XR_ONES_LIKE_TYPE_SETTING(EK_GRIB_DATA) + # Will return a dataarray because that is first value in type-set Union + assert isinstance(ek_reader_result, xr.DataArray) + assert ek_reader_result.equals(xr.ones_like(EK_GRIB_DATA.to_xarray()["2t"])) + + +def test_transform_module_inputs_gribdata_to_xarray(): + # Check EK GribData object + EK_GRIB_DATA = from_source("file", earthkit_test_data_file("test_single.grib")) + ek_reader_result = WRAPPED_DUMMY_MODULE.xarray_ones_like(EK_GRIB_DATA) + # Data array because type-setting of function has dataarray first + assert isinstance(ek_reader_result, xr.DataArray) + assert ek_reader_result.equals(xr.ones_like(EK_GRIB_DATA.to_xarray()["2t"])) + + +def test_transform_function_inputs_wrapper_to_xarray(): + # EK XarrayWrapper object + ek_wrapper_result = WRAPPED_XR_ONES_LIKE(EK_XARRAY_WRAPPER) + assert isinstance(ek_wrapper_result, xr.Dataset) + assert ek_wrapper_result.equals(xr.ones_like(EK_XARRAY_WRAPPER.to_xarray())) + # EK NumpyWrapper object + ek_wrapper_result = WRAPPED_XR_ONES_LIKE(EK_NUMPY_WRAPPER) + assert isinstance(ek_wrapper_result, xr.DataArray) + assert ek_wrapper_result.equals(xr.ones_like(TEST_DA)) + + +def test_transform_module_inputs_wrapper_to_xarray(): + # EK XarrayWrapper object + ek_wrapper_result = WRAPPED_DUMMY_MODULE.xarray_ones_like(EK_XARRAY_WRAPPER) + # Will return a dataarray because that is first value in type-set Union + assert isinstance(ek_wrapper_result, xr.DataArray) + assert ek_wrapper_result.equals(xr.ones_like(EK_XARRAY_WRAPPER.to_xarray()["test"])) + # EK NumpyWrapper object + ek_wrapper_result = WRAPPED_DUMMY_MODULE.xarray_ones_like(EK_NUMPY_WRAPPER) + assert isinstance(ek_wrapper_result, xr.DataArray) + assert ek_wrapper_result.equals(xr.ones_like(TEST_DA)) + + +def test_transform_function_inputs_gribdata_to_numpy(): + # Test with Earthkit.data GribData object + EK_GRIB_DATA = from_source("file", earthkit_test_data_file("test_single.grib")) + assert WRAPPED_NP_MEAN(EK_GRIB_DATA) == np.mean(EK_GRIB_DATA.to_numpy()) + assert isinstance(WRAPPED_NP_MEAN(EK_GRIB_DATA), np.float64) + + +def test_transform_function_inputs_gribdata_to_numpy_typesetting(): + # Test with Earthkit.data GribData object + EK_GRIB_DATA = from_source("file", earthkit_test_data_file("test_single.grib")) + result = WRAPPED_NP_MEAN_TYPE_SETTING(EK_GRIB_DATA) + assert result == np.mean(EK_GRIB_DATA.to_numpy()) + assert isinstance(result, np.float64) + + +def test_transform_module_inputs_gribdata_to_numpy(): + # Test with Earthkit.data GribData object + EK_GRIB_DATA = from_source("file", earthkit_test_data_file("test_single.grib")) + result = WRAPPED_DUMMY_MODULE.numpy_mean(EK_GRIB_DATA) + assert result == np.mean(EK_GRIB_DATA.to_numpy()) + assert isinstance(result, np.float64) + + +def test_transform_function_inputs_wrapper_to_numpy(): + # Test with Earthkit.data XarrayWrapper object + ek_object_result = WRAPPED_NP_MEAN(EK_XARRAY_WRAPPER) + assert ek_object_result == np.mean(TEST_DS) + # assert isinstance(ek_object_result, type(EK_XARRAY_WRAPPER.data)) + + # Test with Earthkit.data NumpyWrapper object + ek_object_result = WRAPPED_NP_MEAN(EK_NUMPY_WRAPPER) + assert ek_object_result == np.mean(TEST_NP) + assert isinstance(ek_object_result, np.float64) + + +def test_transform_module_inputs_wrapper_to_numpy(): + # Test with Earthkit.data XarrayWrapper object + ek_object_result = WRAPPED_DUMMY_MODULE.numpy_mean(EK_XARRAY_WRAPPER) + assert ek_object_result == np.mean(TEST_DS) + # TODO: Without definition of convert_types, WRAPPERS are converted to first arguement + # This could be addressed in future versions + # assert isinstance(ek_object_result, type(EK_XARRAY_WRAPPER.data)) + + # Test with Earthkit.data NumpyWrapper object + ek_object_result = WRAPPED_DUMMY_MODULE.numpy_mean(EK_NUMPY_WRAPPER) + assert ek_object_result == np.mean(TEST_NP) + assert isinstance(ek_object_result, np.float64) + + +def test_transform_function_inputs_xarray_to_numpy(): + # Test with xarray.DataArray object + assert WRAPPED_NP_MEAN(TEST_DA) == np.mean(TEST_NP) + assert WRAPPED_NP_MEAN(TEST_DA) == np.mean(TEST_DA) + assert isinstance(WRAPPED_NP_MEAN(TEST_DA), xr.DataArray) + + +def test_transform_function_inputs_pandas_to_numpy(): + # Test with pandas.DataFrame object, axis=0 is to preserve the pandas.DataFrame, this is pandas syntax + result = WRAPPED_NP_MEAN(TEST_DF, axis=0) + assert result.equals(np.mean(TEST_DF, axis=0)) + assert isinstance(result, pd.core.series.Series) + + # # Test without axis declartion, pandas>2 returns a scalar mean of the dataframe + # # Omitted as breaking change in pandas makes tests unstable + # result = WRAPPED_NP_MEAN(TEST_DF) + # assert result == np.mean(TEST_DF) + # assert isinstance(result, np.float64)