Skip to content

Commit

Permalink
Convert TableSourceWrappers to DataFrameWrapper on injection
Browse files Browse the repository at this point in the history
This way users never have to bother with converting TableSourceWrapper
on their own, something that can be annoying for workflows.
  • Loading branch information
jiffyclub committed Jul 30, 2014
1 parent ec56452 commit a6f9792
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 5 deletions.
11 changes: 7 additions & 4 deletions urbansim/sim/simulation.py
Expand Up @@ -449,10 +449,13 @@ def _collect_injectables(names):
'not all injectables found. '
'missing: {}'.format(names - set(dicts.keys())))

return {
name: thing
if not isinstance(thing, _InjectableFuncWrapper) else thing()
for name, thing in dicts.items()}
for name, thing in dicts.items():
if isinstance(thing, _InjectableFuncWrapper):
dicts[name] = thing()
elif isinstance(thing, TableSourceWrapper):
dicts[name] = thing.convert()

return dicts


def add_table(table_name, table):
Expand Down
8 changes: 7 additions & 1 deletion urbansim/sim/tests/test_simulation.py
Expand Up @@ -230,13 +230,19 @@ def zzz():
def injected():
return 'injected'

@sim.table_source('source')
def source():
return df

with pytest.raises(KeyError):
sim._collect_injectables(['asdf'])

names = ['df', 'df_func', 'answer', 'injected']
names = ['df', 'df_func', 'answer', 'injected', 'source']
things = sim._collect_injectables(names)

assert set(things.keys()) == set(names)
assert isinstance(things['source'], sim.DataFrameWrapper)
pdt.assert_frame_equal(things['source']._frame, df)


def test_injectables(clear_sim):
Expand Down

0 comments on commit a6f9792

Please sign in to comment.