fix float cast to int/id in expressions

scrunch/expressions.py

@@ -46,6 +46,7 @@
 import six
 
 import scrunch
+from scrunch.helpers import is_number
 from scrunch.variables import validate_variable_url
 
 import sys
@@ -540,7 +541,7 @@ def category_ids(var_id, var_value, variables=variables):
                 # {'values': [val1, val2, ...]}
                 value = []
                 for val in var_value:
-                    if str(val).isdigit():
+                    if is_number(val):
                         # val1 is an id already
                         value.append(val)
                         continue
@@ -635,11 +636,14 @@ def _process(obj, variables):
                             values.append(subitem)
                     subitems.append(subitem)
 
-                has_value = any('value' in item for item in subitems
-                    if not str(item).isdigit())
+                has_value = any(
+                    'value' in item for item in subitems if not is_number(item)
+                )
+
+                has_variable = any(
+                    'variable' in item for item in subitems if not is_number(item)
+                )
 
-                has_variable = any('variable' in item for item in subitems
-                    if not str(item).isdigit())
                 if has_value and has_variable:
                     subitems, needs_wrap = ensure_category_ids(subitems)
                 obj[key] = subitems

scrunch/helpers.py

@@ -320,3 +320,11 @@ def generate_subvariable_codes(dataset, subvariables):
         existing_aliases.add(unique_code)
         subvariable_codes.append(unique_code)
     return subvariable_codes
+
+
+def is_number(value):
+    try:
+        float(value)
+        return True
+    except (ValueError, TypeError):
+        return False

scrunch/tests/test_expressions.py

@@ -552,6 +552,38 @@ def test_parse_value_in_list(self):
             ]
         }
 
+    def test_parse_float_value_in_list(self):
+        expr = "country_cat in [1.0]"
+        expected = {
+            'function': 'in',
+            'args': [
+                {
+                    'variable': 'country_cat'
+                },
+                {
+                    'value': [1.0]
+                }
+            ]
+        }
+        expr_obj = parse_expr(expr)
+        assert expr_obj == expected
+
+    def test_parse_integer_value_in_list(self):
+        expr = "country_cat in [1]"
+        expected = {
+            'function': 'in',
+            'args': [
+                {
+                    'variable': 'country_cat'
+                },
+                {
+                    'value': [1]
+                }
+            ]
+        }
+        expr_obj = parse_expr(expr)
+        assert expr_obj == expected
+
     def test_r_in(self):
         expr = "q1 in [1, 2, r(4,7), r(10, 12)]"
         expected_expr_obj = {
@@ -2623,6 +2655,91 @@ def test_multiple_response_any(self, adapter):
             ]
         }
 
+    def test_in_expression_list_integer(self):
+        var_id = '0001'
+        var_alias = 'country_cat'
+        var_type = 'categorical'
+        var_url = '{}variables/{}/'.format(self.ds_url, var_id)
+        categories = [
+            {
+                'name': 'argentina',
+                'id': 1
+            },
+            {
+                'name': 'australia',
+                'id': 2
+            },
+        ]
+
+        table_mock = mock.MagicMock(metadata={
+            var_id: {
+                'id': var_id,
+                'alias': var_alias,
+                'type': var_type,
+                'categories': categories,
+            }
+        })
+        ds = mock.MagicMock()
+        ds.self = self.ds_url
+        ds.follow.return_value = table_mock
+
+        expr = "country_cat in [1]"
+        expr_obj = process_expr(parse_expr(expr), ds)
+        assert expr_obj == {
+            'function': 'in',
+            'args': [
+                {
+                    'variable': var_url
+                },
+                {
+                    'value': [1]
+                }
+            ]
+        }
+
+    def test_in_expression_list_floats(self):
+        var_id = '0001'
+        var_alias = 'country_cat'
+        var_type = 'categorical'
+        var_url = '{}variables/{}/'.format(self.ds_url, var_id)
+        categories = [
+            {
+                'name': 'argentina',
+                'id': 1
+            },
+            {
+                'name': 'australia',
+                'id': 2
+            },
+        ]
+
+        table_mock = mock.MagicMock(metadata={
+            var_id: {
+                'id': var_id,
+                'alias': var_alias,
+                'type': var_type,
+                'categories': categories,
+            }
+        })
+        ds = mock.MagicMock()
+        ds.self = self.ds_url
+        ds.follow.return_value = table_mock
+
+        expr = "country_cat in [1.0]"
+        parsed_expr = parse_expr(expr)
+        expr_obj = process_expr(parsed_expr, ds)
+        assert expr_obj == {
+            'function': 'in',
+            'args': [
+                {
+                    'variable': var_url
+                },
+                {
+                    'value': [1.0]
+                }
+            ]
+        }
+
 
 class TestExpressionPrettify(TestCase):