def test_build_wrapper_function():
def func(a, b):
pass
# Actual function inputs & DS2 variables current dont have to match.
result = build_wrapper_function(func, [DS2Variable('a', 'int', False),
DS2Variable('x', 'float', True)],
array_input=False,
return_msg=False)
assert isinstance(result, str)
assert '"Output: x"\n' in result
def test_wrapper_renamed():
"""Check wrap_predict_method output with custom name."""
target = """
def renamed_wrapper(a, b):
"Output: c, msg"
result = None
msg = None
try:
global _compile_error
if _compile_error is not None:
raise _compile_error
import numpy as np
import pandas as pd
if a is None: a = np.nan
if b is None: b = np.nan
input_array = np.array([a, b]).reshape((1, -1))
columns = ["a", "b"]
input_df = pd.DataFrame(data=input_array, columns=columns)
result = dummy_func(input_df)
result = tuple(result.ravel()) if hasattr(result, "ravel") else tuple(result)
if len(result) == 0:
result = tuple(None for i in range(1))
elif "numpy" in str(type(result[0])):
result = tuple(np.asscalar(i) for i in result)
except Exception as e:
from traceback import format_exc
msg = str(e) + format_exc()
if result is None:
result = tuple(None for i in range(1))
return result + (msg, )
""".rstrip()
code = build_wrapper_function(dummy_func, [
DS2Variable('a', float, False),
DS2Variable('b', float, False),
DS2Variable('c', float, True)
],
array_input=True,
name='renamed_wrapper')
assert code == target