def test_let_elementwise():
df = pd.DataFrame()
df["A"] = ["hello", "alice", "bob"]
w = Wrap(df)
def first(s):
return s[0]
w = w.let_elementwise(first1=first)
wnew = w.extend("B = first1(A)")
assert ["A", "B"] == list(wnew.df.columns)
assert ["h", "a", "b"] == list(wnew.df["B"])
def test_let_elementwise_multiple_params():
df = pd.DataFrame()
df["A"] = [1, 2]
w = Wrap(df)
def x(a, b):
# add assertion to ensure the function is called elementwise
assert a == 1 or a == 2
return a + b
w = w.let_elementwise(x=x)
wnew = w.extend("B = x(A, A + 1)")
assert [3, 5] == list(wnew.df["B"])