diff --git a/scrunch/expressions.py b/scrunch/expressions.py index 2b3e4084..eafa5407 100644 --- a/scrunch/expressions.py +++ b/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 diff --git a/scrunch/helpers.py b/scrunch/helpers.py index 5ee921dc..1990e54e 100644 --- a/scrunch/helpers.py +++ b/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 diff --git a/scrunch/tests/test_expressions.py b/scrunch/tests/test_expressions.py index 26b92a1e..135094fb 100644 --- a/scrunch/tests/test_expressions.py +++ b/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):