From 6d6a4a0bb41eb496085b2114b05ff29ff1920233 Mon Sep 17 00:00:00 2001 From: Matt Davis Date: Wed, 14 May 2014 10:56:12 -0700 Subject: [PATCH] Return indexes of added rows from TransitionModel.transition. --- urbansim/models/tests/test_transition.py | 5 ++++- urbansim/models/transition.py | 4 +++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/urbansim/models/tests/test_transition.py b/urbansim/models/tests/test_transition.py index 528d1d15..25d06f50 100644 --- a/urbansim/models/tests/test_transition.py +++ b/urbansim/models/tests/test_transition.py @@ -290,7 +290,7 @@ def test_transition_model(basic_df, grow_targets_filters, totals_col, year): {'z': ['a', 'b', 'c', 'd', 'e'], 'thing_id': basic_df.index}) - new, new_linked = model.transition( + new, added, new_linked = model.transition( basic_df, year, linked_tables={'linked': (linked_table, 'thing_id')}) assert len(new) == grow_targets_filters[totals_col].sum() @@ -298,3 +298,6 @@ def test_transition_model(basic_df, grow_targets_filters, totals_col, year): assert len(new_linked['linked']) == grow_targets_filters[totals_col].sum() assert new.index.values.max() in new_linked['linked'].thing_id.values assert new_linked['linked'].index.values.max() == 5 + assert added.isin(new.index).all() + assert not added.isin(basic_df.index).any() + npt.assert_array_equal(added.values, [basic_df.index.values.max() + 1]) diff --git a/urbansim/models/transition.py b/urbansim/models/transition.py index 166a611c..00ad00d0 100644 --- a/urbansim/models/transition.py +++ b/urbansim/models/transition.py @@ -444,6 +444,8 @@ def transition(self, data, year, linked_tables=None): ------- updated : pandas.DataFrame Table with rows removed or added. + added : pandas.Series + Indexes of new rows in `updated`. updated_links : dict of pandas.DataFrame """ @@ -456,4 +458,4 @@ def transition(self, data, year, linked_tables=None): updated_links[table_name] = \ _update_linked_table(table, col, added, copied, removed) - return updated, updated_links + return updated, added, updated_links