Function for mapping column names to the tables that contain them.

urbansim/utils/misc.py

@@ -5,8 +5,10 @@
 from __future__ import print_function
 
 import os
+
 import numpy as np
 import pandas as pd
+import toolz
 
 
 def _mkifnotexists(folder):
@@ -285,3 +287,34 @@ def pandasdfsummarytojson(df, ndigits=3):
     """
     df = df.transpose()
     return {k: _pandassummarytojson(v, ndigits) for k, v in df.iterrows()}
+
+
+def column_map(tables, columns):
+    """
+    Take a list of tables and a list of column names and resolve which
+    columns come from which table.
+
+    Parameters
+    ----------
+    tables : sequence of _DataFrameWrapper or _TableFuncWrapper
+        Could also be sequence of modified pandas.DataFrames, the important
+        thing is that they have ``.name`` and ``.columns`` attributes.
+    columns : sequence of str
+        The column names of interest.
+
+    Returns
+    -------
+    col_map : dict
+        Maps table names to lists of column names.
+
+    """
+    if not columns:
+        return {t.name: None for t in tables}
+
+    columns = set(columns)
+    colmap = {t.name: list(set(t.columns).intersection(columns)) for t in tables}
+    foundcols = toolz.reduce(lambda x, y: x.union(y), (set(v) for v in colmap.values()))
+    if foundcols != columns:
+        raise RuntimeError('Not all required columns were found. '
+                           'Missing: {}'.format(list(columns - foundcols)))
+    return colmap

urbansim/utils/tests/test_misc.py

@@ -0,0 +1,35 @@
+import pytest
+
+from .. import misc
+
+
+class _FakeTable(object):
+    def __init__(self, name, columns):
+        self.name = name
+        self.columns = columns
+
+
+@pytest.fixture
+def fta():
+    return _FakeTable('a', ['aa', 'ab', 'ac'])
+
+
+@pytest.fixture
+def ftb():
+    return _FakeTable('b', ['bx', 'by', 'bz'])
+
+
+def test_column_map_raises(fta, ftb):
+    with pytest.raises(RuntimeError):
+        misc.column_map([fta, ftb], ['aa', 'by', 'bz', 'cw'])
+
+
+def test_column_map_none(fta, ftb):
+    assert misc.column_map([fta, ftb], None) == {'a': None, 'b': None}
+
+
+def test_column_map(fta, ftb):
+    assert misc.column_map([fta, ftb], ['aa', 'by', 'bz']) == \
+        {'a': ['aa'], 'b': ['by', 'bz']}
+    assert misc.column_map([fta, ftb], ['by', 'bz']) == \
+        {'a': [], 'b': ['by', 'bz']}