Merge pull request #39 from synthicity/query-str

Support passing strings to apply_filter_query

urbansim/models/tests/test_util.py

@@ -71,6 +71,16 @@ def test_apply_filter_query_no_filter(test_df):
     pdt.assert_frame_equal(filtered, expected)
 
 
+def test_apply_filter_query_str(test_df):
+    filters = 'col1 < 3'
+    filtered = util.apply_filter_query(test_df, filters)
+    expected = pd.DataFrame(
+        {'col1': [0, 1, 2],
+         'col2': [5, 6, 7]},
+        index=['a', 'b', 'c'])
+    pdt.assert_frame_equal(filtered, expected)
+
+
 @pytest.mark.parametrize('name, val, filter_exp', [
     ('x', 1, 'x == 1'),
     ('x', 'a', "x == 'a'"),

urbansim/models/util.py

@@ -18,9 +18,10 @@ def apply_filter_query(df, filters=None):
     Parameters
     ----------
     df : pandas.DataFrame
-    filters : list of str, optional
+    filters : list of str or str, optional
         List of filters to apply. Will be joined together with
-        ' and ' and passed to DataFrame.query.
+        ' and ' and passed to DataFrame.query. A string will be passed
+        straight to DataFrame.query.
         If not supplied no filtering will be done.
 
     Returns
@@ -29,7 +30,10 @@ def apply_filter_query(df, filters=None):
 
     """
     if filters:
-        query = ' and '.join(filters)
+        if isinstance(filters, str):
+            query = filters
+        else:
+            query = ' and '.join(filters)
         return df.query(query)
     else:
         return df