Add .to_frame method to Account

urbansim/accounts.py

@@ -4,9 +4,33 @@
 """
 from collections import namedtuple
 
+import pandas as pd
+import toolz
+
 
 Transaction = namedtuple('Transaction', ('amount', 'subaccount', 'metadata'))
 
+# column names that are always present in DataFrames of transactions
+COLS = ['amount', 'subaccount']
+
+
+def _column_names_from_metadata(dicts):
+    """
+    Get the unique set of keys from a list of dictionaries.
+
+    Parameters
+    ----------
+    dicts : iterable
+        Sequence of dictionaries.
+
+    Returns
+    -------
+    keys : list
+        Unique set of keys.
+
+    """
+    return list(toolz.unique(toolz.concat(dicts)))
+
 
 class Account(object):
     """
@@ -38,11 +62,11 @@ def add_transaction(self, amount, subaccount=None, metadata=None):
         metadata : dict, optional
             Any extra metadata to record with the transaction.
             (E.g. Info about where the money is coming from or going.)
+            May not contain keys 'amount' or 'subaccount'.
 
         """
         metadata = metadata or {}
-        t = Transaction(amount, subaccount, metadata)
-        self.transactions.append(t)
+        self.transactions.append(Transaction(amount, subaccount, metadata))
         self.balance += amount
 
     def add_transactions(self, transactions):
@@ -58,3 +82,19 @@ def add_transactions(self, transactions):
         """
         for t in transactions:
             self.add_transaction(*t)
+
+    def to_frame(self):
+        """
+        Return transactions as a pandas DataFrame.
+
+        """
+        col_names = _column_names_from_metadata(
+            t.metadata for t in self.transactions)
+
+        trow = lambda t: (
+            toolz.concatv(
+                (t.amount, t.subaccount),
+                (t.metadata.get(c) for c in col_names)))
+        rows = [trow(t) for t in self.transactions]
+
+        return pd.DataFrame(rows, columns=COLS + col_names)

urbansim/tests/test_accounts.py

@@ -1,4 +1,6 @@
+import pandas as pd
 import pytest
+from pandas.util import testing as pdt
 
 from .. import accounts
 
@@ -47,3 +49,26 @@ def test_add_transactions(acc, acc_bal):
 
     assert len(acc.transactions) == 2
     assert acc.balance == acc_bal + t1[0] + t2[0]
+
+
+def test_column_names_from_metadata():
+    cnfm = accounts._column_names_from_metadata
+
+    assert cnfm([]) == []
+    assert cnfm([{'a': 1, 'b': 2}]) == ['a', 'b']
+    assert cnfm([{'a': 1}, {'b': 2}]) == ['a', 'b']
+    assert cnfm([{'a': 1, 'b': 2}, {'a': 3, 'b': 4}]) == ['a', 'b']
+
+
+def test_to_frame(acc, acc_bal):
+    t1 = accounts.Transaction(200, ('a', 'b', 'c'), None)
+    t2 = (-50, None, {'to': 'Acme Corp.'})
+    acc.add_transactions((t1, t2))
+
+    expected = pd.DataFrame(
+        [[200, ('a', 'b', 'c'), None],
+         [-50, None, 'Acme Corp.']],
+        columns=['amount', 'subaccount', 'to'])
+
+    df = acc.to_frame()
+    pdt.assert_frame_equal(df, expected)