Add "not" to reserved words in columns_in_filters.

UDST · Jul 16, 2014 · 2c1cce4 · 2c1cce4
1 parent a0c9cb0
commit 2c1cce4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/urbansim/models/tests/test_util.py b/urbansim/models/tests/test_util.py
@@ -166,6 +166,9 @@ def test_columns_in_filters():
     filters = 'a > 5 and 10 < b < 20 and c == "x" and (d > 20 or e < 50)'
     assert set(util.columns_in_filters(filters)) == {'a', 'b', 'c', 'd', 'e'}
 
+    filters = 'a in [1, 2, 3] or b not in ["x", "y", "z"]'
+    assert set(util.columns_in_filters(filters)) == {'a', 'b'}
+
 
 def test_columns_in_formula():
     formula = 'a ~ b + c - d * e:f'

diff --git a/urbansim/models/util.py b/urbansim/models/util.py
@@ -273,13 +273,13 @@ def columns_in_filters(filters):
         filters = ' '.join(filters)
 
     columns = []
-    reserved = {'and', 'or', 'in'}
+    reserved = {'and', 'or', 'in', 'not'}
 
     for toknum, tokval, _, _, _ in generate_tokens(StringIO(filters).readline):
         if toknum == NAME and tokval not in reserved:
             columns.append(tokval)
 
-    return list(set(columns))
+    return list(toolz.unique(columns))
 
 
 def _tokens_from_patsy(node):
@@ -337,4 +337,4 @@ def columns_in_formula(formula):
                 if toknum == NAME:
                     columns.append(tokval)
 
-    return list(set(columns))
+    return list(toolz.unique(columns))