def test_boolean_support(engine, tsh):
@func('op-with-boolean-kw')
def customseries(zeroes: bool = False) -> pd.Series:
return pd.Series(np.array([1.0, 2.0, 3.0]) * zeroes,
index=pd.date_range(dt(2019, 1, 1),
periods=3,
freq='D'))
tsh.register_formula(engine, 'no-zeroes', '(op-with-boolean-kw)')
tsh.register_formula(engine, 'zeroes', '(op-with-boolean-kw #:zeroes #t)')
ts1 = tsh.get(engine, 'no-zeroes')
assert_df("""
2019-01-01 0.0
2019-01-02 0.0
2019-01-03 0.0
""", ts1)
ts2 = tsh.get(engine, 'zeroes')
assert_df("""
2019-01-01 1.0
2019-01-02 2.0
2019-01-03 3.0
""", ts2)
FUNCS.pop('op-with-boolean-kw')
def test_custom_metadata(engine, tsh):
@func('customseries')
def customseries() -> pd.Series:
return pd.Series([1.0, 2.0, 3.0],
index=pd.date_range(dt(2019, 1, 1),
periods=3,
freq='D'))
@finder('customseries')
def find(cn, tsh, tree):
return {
tree[0]: {
'index_type': 'datetime64[ns]',
'index_dtype': '|M8[ns]',
'tzaware': False,
'value_type': 'float64',
'value_dtype': '<f8'
}
}
tsh.register_formula(engine, 'custom', '(+ 3 (customseries))', False)
meta = tsh.metadata(engine, 'custom')
assert meta == {
'index_type': 'datetime64[ns]',
'index_dtype': '|M8[ns]',
'tzaware': False,
'value_type': 'float64',
'value_dtype': '<f8'
}
# cleanup
FUNCS.pop('customseries')
def test_staircase(engine, tsh):
tsh.register_formula(engine, 's-addition',
'(add (series "sa") (series "sb"))', False)
for day in (1, 2, 3, 4, 5):
idate = utcdt(2018, 1, day)
for name in 'ab':
ts = pd.Series([day / 2.] * 5,
index=pd.date_range(dt(2018, 1, day),
periods=5,
freq='D'))
tsh.update(engine, ts, 's' + name, 'Babar', insertion_date=idate)
ts = tsh.staircase(engine, 's-addition', delta=pd.Timedelta(hours=12))
assert_df(
"""
2018-01-02 1.0
2018-01-03 2.0
2018-01-04 3.0
2018-01-05 4.0
2018-01-06 5.0
2018-01-07 5.0
2018-01-08 5.0
2018-01-09 5.0
""", ts)
# this is not allowed in the staircase fast-path
# hence we will take the slow path
@func('identity')
def identity(series: pd.Series) -> pd.Series:
return series
tsh.register_formula(engine, 'slow-down', '(identity (series "sa"))',
False)
tsh.register_formula(engine, 's-addition-not-fast',
'(add (series "slow-down") (series "sb"))', False)
ts = tsh.staircase(engine,
's-addition-not-fast',
delta=pd.Timedelta(hours=12))
assert_df(
"""
2018-01-02 1.0
2018-01-03 2.0
2018-01-04 3.0
2018-01-05 4.0
2018-01-06 5.0
2018-01-07 5.0
2018-01-08 5.0
2018-01-09 5.0
""", ts)
# cleanup
FUNCS.pop('identity')
def find_operators(self, cn, tree):
ops = {tree[0]: FUNCS.get(tree[0])}
for item in tree:
if isinstance(item, list):
newops = self.find_operators(cn, item)
ops.update(newops)
return ops
def test_new_func(engine, tsh):
@func('identity')
def identity(series: pd.Series) -> pd.Series:
return series
tsh.register_formula(engine, 'identity', '(identity (series "id-a"))',
False)
ts = pd.Series([1, 2, 3],
index=pd.date_range(dt(2019, 1, 1), periods=3, freq='D'))
tsh.update(engine, ts, 'id-a', 'Babar')
ts = tsh.get(engine, 'identity')
assert_df("""
2019-01-01 1.0
2019-01-02 2.0
2019-01-03 3.0
""", ts)
# cleanup
FUNCS.pop('identity')
def test_options(engine, tsh):
@func('dummy')
def dummy(option: int = None) -> pd.Series:
series = pd.Series([1, 2, 3],
index=pd.date_range(dt(2019, 1, 1),
periods=3,
freq='D'))
series.options = {'option': option}
return series
tsh.register_formula(engine, 'test_options', '(* 3 (dummy #:option 42))',
False)
ts = tsh.get(engine, 'test_options')
assert_df("""
2019-01-01 3
2019-01-02 6
2019-01-03 9
""", ts)
assert ts.options == {'option': 42}
FUNCS.pop('dummy')
def test_ifunc(engine, tsh):
@func('shifted')
def shifted(__interpreter__, name: str, days: int = 0) -> pd.Series:
args = __interpreter__.getargs.copy()
fromdate = args.get('from_value_date')
todate = args.get('to_value_date')
if fromdate:
args['from_value_date'] = fromdate + timedelta(days=days)
if todate:
args['to_value_date'] = todate + timedelta(days=days)
return __interpreter__.get(name, args)
@finder('shifted')
def find_series(cn, tsh, stree):
return {stree[1]: tsh.metadata(cn, stree[1])}
tsh.register_formula(engine, 'shifting',
'(+ 0 (shifted "shiftme" #:days -1))', False)
ts = pd.Series([1, 2, 3, 4, 5],
index=pd.date_range(dt(2019, 1, 1), periods=5, freq='D'))
tsh.update(engine,
ts,
'shiftme',
'Babar',
insertion_date=utcdt(2019, 1, 1))
ts = tsh.get(engine, 'shifting')
assert_df(
"""
2019-01-01 1.0
2019-01-02 2.0
2019-01-03 3.0
2019-01-04 4.0
2019-01-05 5.0
""", ts)
ts = tsh.get(engine,
'shifting',
from_value_date=dt(2019, 1, 3),
to_value_date=dt(2019, 1, 4))
assert_df("""
2019-01-02 2.0
2019-01-03 3.0
""", ts)
# now, history
ts = pd.Series([1, 2, 3, 4, 5],
index=pd.date_range(dt(2019, 1, 2), periods=5, freq='D'))
tsh.update(engine,
ts,
'shiftme',
'Babar',
insertion_date=utcdt(2019, 1, 2))
hist = tsh.history(engine, 'shifting')
assert_hist(
"""
insertion_date value_date
2019-01-01 00:00:00+00:00 2019-01-01 1.0
2019-01-02 2.0
2019-01-03 3.0
2019-01-04 4.0
2019-01-05 5.0
2019-01-02 00:00:00+00:00 2019-01-01 1.0
2019-01-02 1.0
2019-01-03 2.0
2019-01-04 3.0
2019-01-05 4.0
2019-01-06 5.0
""", hist)
hist = tsh.history(engine,
'shifting',
from_value_date=dt(2019, 1, 3),
to_value_date=dt(2019, 1, 4))
assert_hist(
"""
insertion_date value_date
2019-01-01 00:00:00+00:00 2019-01-03 3.0
2019-01-04 4.0
2019-01-02 00:00:00+00:00 2019-01-03 2.0
2019-01-04 3.0
""", hist)
# cleanup
FUNCS.pop('shifted')
def test_editor_new_operator(mapi):
@func('genrandomseries')
def genrandomseries() -> pd.Series:
return pd.Series([1.0, 2.0, 3.0],
index=pd.date_range(dt(2019, 1, 1),
periods=3,
freq='D'))
@finder('genrandomseries')
def genrandomseries(cn, tsh, tree):
return {
tree[0]: {
'index_dtype': '<M8[ns]',
'index_type': 'datetime64[ns]',
'tzaware': False,
'value_dtype': '<f8',
'value_type': 'float64'
}
}
@func('frobulated')
def frobulate(a: str, b: str) -> pd.Series:
sa = mapi.get(a)
sb = mapi.get(b)
return (sa + 1) * sb
mapi.register_formula(
'random',
'(genrandomseries)',
)
ts = mapi.get('random')
assert_df("""
2019-01-01 1.0
2019-01-02 2.0
2019-01-03 3.0
""", ts)
presenter = fancypresenter(mapi, 'random', {})
info = [{k: v
for k, v in info.items() if k != 'ts'}
for info in presenter.infos]
assert info == [{
'name': 'random',
'type': 'formula'
}, {
'name': 'genrandomseries',
'type': 'primary'
}]
mapi.update(
'new-op',
pd.Series([1, 2, 3],
index=pd.date_range(dt(2019, 1, 1), periods=3, freq='D')),
'Baabar')
mapi.register_formula(
'frobulating', '(* 2 (add (series "random")'
' (frobulated "new-op" "new-op")))')
presenter = fancypresenter(mapi, 'frobulating', {})
info = [{k: v
for k, v in info.items() if k != 'ts'}
for info in presenter.infos]
assert info == [{
'name': 'frobulating',
'type': 'formula'
}, {
'name': 'random',
'type': 'formula'
}]
# cleanup
FUNCS.pop('genrandomseries')
FUNCS.pop('frobulated')