Permalink
Browse files

clarify run output tables: update args and add test

  • Loading branch information...
bridwell committed Jul 1, 2016
1 parent 4a714d6 commit ef03a5d6b3043b8f16af4011f4169da9a0d71cee
Showing with 34 additions and 9 deletions.
  1. +8 −9 orca/orca.py
  2. +26 −0 orca/tests/test_orca.py
View
@@ -1841,11 +1841,10 @@ def write_tables(fname, table_names=None, prefix=None):
with pd.get_store(fname, mode='a') as store:
for t in tables:
store[key_template.format(t.name)] = t.to_frame()
store.close()
def run(steps, iter_vars=None, data_out=None, out_interval=1,
base_tables=None, run_tables=None):
out_base_tables=None, out_run_tables=None):
"""
Run steps in series, optionally repeatedly over some sequence.
The current iteration variable is set as a global injectable
@@ -1884,18 +1883,18 @@ def run(steps, iter_vars=None, data_out=None, out_interval=1,
max_i = len(iter_vars)
# get the tables to write out
if base_tables is None or run_tables is None:
if out_base_tables is None or out_run_tables is None:
step_tables = get_step_table_names(steps)
if base_tables is None:
base_tables = step_tables
if out_base_tables is None:
out_base_tables = step_tables
if run_tables is None:
run_tables = step_tables
if out_run_tables is None:
out_run_tables = step_tables
# write out the base (inputs)
if data_out:
write_tables(data_out, base_tables, 'base')
write_tables(data_out, out_base_tables, 'base')
# run the steps
for i, var in enumerate(iter_vars, start=1):
@@ -1929,7 +1928,7 @@ def run(steps, iter_vars=None, data_out=None, out_interval=1,
# write out the results for the current iteration
if data_out:
if (i - 1) % out_interval == 0 or i == max_i:
write_tables(data_out, run_tables, var)
write_tables(data_out, out_run_tables, var)
clear_cache(scope=_CS_ITER)
View
@@ -955,6 +955,32 @@ def step(iter_var, table):
store['10/table'][year_key(x)], series_year(x))
def test_run_and_write_tables_out_tables_provided(df, store_name):
table_names = ['table', 'table2', 'table3']
for t in table_names:
orca.add_table(t, df)
@orca.step()
def step(iter_var, table, table2):
return
orca.run(
['step'],
iter_vars=range(1),
data_out=store_name,
out_base_tables=table_names,
out_run_tables=['table'])
with pd.get_store(store_name, mode='r') as store:
for t in table_names:
assert 'base/{}'.format(t) in store
assert '0/table' in store
assert '0/table2' not in store
assert '0/table3' not in store
def test_get_raw_table(df):
orca.add_table('table1', df)

0 comments on commit ef03a5d

Please sign in to comment.