WIP

Author: slobodan-ilic

scrunch/datasets.py

@@ -1129,7 +1129,7 @@ def create_fill_values(self, variables, name, alias, description=''):
         if "variable" in else_case and "name" in else_case:
             raise ValueError("Else case can be either variable or category not both")
 
-        aliases = {c["variable"] for c in variables}
+        aliases = {c["var"] for c in variables}
         vars_by_alias = self.resource.variables.by("alias")
         types = {vars_by_alias[al]["type"] for al in aliases}
         if types != {"categorical"}:
@@ -1162,13 +1162,13 @@ def create_fill_values(self, variables, name, alias, description=''):
             })
 
         expr = {"function": "case", "args": args}
-        fill_map = {str(cid): {"variable": vars_by_alias[v["variable"]]["id"]}
+        fill_map = {str(cid): {"var": v["var"]}
                     for cid, v in zip(cat_ids, variables)}
 
-        if "variable" in else_case:
+        if "var" in else_case:
             # We are in the case of a default fill, replace the -1 with the new
             # variable
-            fill_map["-1"] = {"variable": vars_by_alias[else_case["variable"]]["id"]}
+            fill_map["-1"] = {"var": else_case["var"]}
 
         fill_expr = {
             "function": "fill",

scrunch/expressions.py

@@ -51,7 +51,7 @@
 
 import sys
 
-PY311 = sys.version_info[:2] == (3, 11)
+GT_PY_311 = sys.version_info[:2] >= (3, 11)
 
 if six.PY2:
     from urllib import urlencode
@@ -292,7 +292,7 @@ def _parse(node, parent=None):
                 # We will take the subvariable alias bit from the subscript
                 # and return an object with the array and subvariable alias
                 array_alias = dict(ast.iter_fields(fields[0][1]))["id"]
-                if PY311:                
+                if GT_PY_311:                
                     name_node = dict(ast.iter_fields(fields[1][1]))
                     subvariable_alias = name_node["id"]
                 else:
@@ -494,17 +494,17 @@ def get_subvariables_resource(var_url, var_index):
     return {sv['alias'].strip('#'): sv['id'] for sv in sub_variables.values()}
 
 
-def _get_categories_from_var_index(var_index, var_url):
-    return var_index[var_url].entity.body.categories
+def _get_categories_from_var_index(vars_by_alias, var_alias):
+    return vars_by_alias[var_alias].entity.body.categories
 
 
-def adapt_multiple_response(var_url, values, var_index):
+def adapt_multiple_response(var_alias, values, vars_by_alias):
     """
     Converts multiple response arguments
     to column.
     :return: the new args for multiple_response
     """
-    aliases = get_subvariables_resource(var_url, var_index)
+    aliases = get_subvariables_resource(var_alias, vars_by_alias)
     result = []
 
     if all(isinstance(value, int) for value in values):
@@ -515,35 +515,38 @@ def adapt_multiple_response(var_url, values, var_index):
         # scenario var.any([subvar1, subvar2])
         # in this scenario, we only want category ids that refers to `selected` categories
         column = [
-            cat.get("id") for cat in _get_categories_from_var_index(var_index, var_url) if cat.get("selected")
+            cat.get("id") for cat in _get_categories_from_var_index(vars_by_alias, var_alias) if cat.get("selected")
         ]
-        variables = [var_id for alias, var_id in aliases.items() if alias in values]
+        subvars = [sva for sva in aliases if sva in values]
 
-    for variable_id in variables:
-        variable_url = "{}subvariables/{}/".format(var_url, variable_id)
+    for sva in subvars:
         result.append({
-            "variable": variable_url,
+            "var": var_alias,
+            'axes': [sva],
             "column": column
         })
 
     return result, True
 
 
-def _update_values_for_multiple_response(new_values, values, subitem, var_index, arrays):
+def _update_values_for_multiple_response(new_values, values, subitem, vars_by_alias, arrays):
     """
     - Multiple response does not need the `value` key, but it relies on the `column` key
     - Remove from `arrays` (subvariable list) the ones that should not be considered
     """
-    var_url = subitem.get("variable", "").split("subvariables")[0]
+    # var_url = subitem.get("variable", "").split("subvariables")[0]
+    var_alias = subitem.get("var")
     column = new_values[0].get("column")
     value = values[0].get("value")
-    if var_url and var_index[var_url]['type'] == 'multiple_response':
+    # if var_url and var_index[var_url]['type'] == 'multiple_response':
+    if var_alias and vars_by_alias[var_alias]['type'] == 'multiple_response':
         if column:
             values[0]['column'] = column
         elif value is not None:
             values[0]['column'] = value
         values[0].pop("value", None)
-        arrays[0] = [new_value["variable"] for new_value in new_values]
+        subvar_ids_by_aliases = {v['alias']: k for k, v in vars_by_alias[var_alias]['entity']['subvariables']['index'].items()}
+        arrays[0] = [subvar_ids_by_aliases[new_value["axes"][0]] for new_value in new_values]
 
 
 def process_expr(obj, ds):
@@ -567,7 +570,7 @@ def ensure_category_ids(subitems, values, arrays, variables=variables):
         def variable_id(variable_url):
             return variable_url.split('/')[-2]
 
-        def category_ids(var_id, var_value, variables=variables):
+        def category_ids(var_alias, var_value, variables=variables):
             value = None
             if isinstance(var_value, list) or isinstance(var_value, tuple):
                 # {'values': [val1, val2, ...]}
@@ -578,7 +581,7 @@ def category_ids(var_id, var_value, variables=variables):
                         value.append(val)
                         continue
                     for var in variables:
-                        if variables[var]['id'] == var_id:
+                        if variables[var]['alias'] == var_alias:
                             if 'categories' in variables[var]:
                                 for cat in variables[var]['categories']:
                                     if cat['name'] == val:
@@ -587,22 +590,21 @@ def category_ids(var_id, var_value, variables=variables):
                                 # variable has no categories, return original
                                 # list of values
                                 value = var_value
-
             elif isinstance(var_value, str):
-                for var in variables:
+                for va, var in variables.items():
                     # if the variable is a date, don't try to process it's categories
-                    if variables[var]['type'] == 'datetime':
+                    if var['type'] == 'datetime':
                         return var_value
-                    if variables[var]['id'] == var_id and 'categories' in variables[var]:
+                    if va == var_alias and 'categories' in var:
                         found = False
-                        for cat in variables[var]['categories']:
+                        for cat in var['categories']:
                             if cat['name'] == var_value:
                                 value = cat['id']
                                 found = True
                                 break
                         if not found:
                             raise ValueError("Couldn't find a category id for category %s in filter for variable %s" % (var_value, var))
-                    elif 'categories' not in variables[var]:
+                    elif 'categories' not in variables[var['alias']]:
                         return var_value
 
             else:
@@ -612,20 +614,19 @@ def category_ids(var_id, var_value, variables=variables):
         # special case for multiple_response variables
         if len(subitems) == 2:
             _variable, _value = subitems
-            var_url = _variable.get('variable')
+            var_alias = _variable.get('var')
             _value_key = next(iter(_value))
-            if _value_key in {'column', "value"} and var_url:
-                if var_url in var_index and var_index[var_url]['type'] == 'multiple_response':
-                    result = adapt_multiple_response(var_url, _value[_value_key], var_index)
+            if _value_key in {'column', "value"} and var_alias:
+                vars_by_alias = {v['alias']: v for _, v in var_index.items()}
+                if var_alias in vars_by_alias and vars_by_alias[var_alias]['type'] == 'multiple_response':
+                    result = adapt_multiple_response(var_alias, _value[_value_key], vars_by_alias)
                     # handle the multiple response type
-                    _update_values_for_multiple_response(result[0], values, subitems[0], var_index, arrays)
+                    _update_values_for_multiple_response(result[0], values, subitems[0], vars_by_alias, arrays)
                     return result
 
         for item in subitems:
-            if isinstance(item, dict) and 'variable' in item and not isinstance(item["variable"], dict):
-                var_id = variable_id(item['variable'])
-            elif isinstance(item, dict) and 'value' in item:
-                item['value'] = category_ids(var_id, item['value'])
+            if isinstance(item, dict) and 'value' in item:
+                item['value'] = category_ids(var_alias, item['value'])
             _subitems.append(item)
 
         return _subitems, True
@@ -657,11 +658,12 @@ def _process(obj, variables):
                         except TypeError:
                             raise ValueError("Invalid variable alias '%s'" % args[0]['variable'])
 
+        new_obj = copy.deepcopy(obj)
         for key, val in obj.items():
             if isinstance(val, dict) and "array" not in val:
                 # This is not an array object, then it's a nested ZCL expression
                 # so we need to proceed for nested processing.
-                obj[key] = _process(val, variables)
+                new_obj[key] = _process(val, variables)
             elif isinstance(val, (list, tuple)):
                 subitems = []
                 for subitem in val:
@@ -671,6 +673,13 @@ def _process(obj, variables):
                             arrays.append(subitem.pop('subvariables'))
                         elif 'value' in subitem or 'column' in subitem:
                             values.append(subitem)
+                        elif 'var' in subitem:
+                            var = variables.get(subitem['var'])
+                            if var['type'] in ARRAY_TYPES and 'axes' not in subitem:
+                                # Add info about the fact that the "var" referenced
+                                # variable is an array, so that the validation can
+                                # work properly in the remaindere of the code.
+                                arrays.append(var['subvariables'])
                     subitems.append(subitem)
 
                 has_value = any(
@@ -683,13 +692,14 @@ def _process(obj, variables):
                     has_value = any('column' 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)
+                    # 'variable' in item for item in subitems if not is_number(item)
+                    'var' in item for item in subitems if not is_number(item)
                 )
 
                 if has_value and has_variable:
                     subitems, needs_wrap = ensure_category_ids(subitems, values, arrays)
 
-                obj[key] = subitems
+                new_obj[key] = subitems
             elif key == 'variable':
                 if isinstance(val, dict) and "array" in val:
                     # This is a subvariable reference with this shape:
@@ -707,7 +717,7 @@ def _process(obj, variables):
                     except KeyError:
                         raise ValueError("Invalid subvariable `%s` for array '%s'" % (subvariables, array_alias))
                     subvar_url = "%svariables/%s/subvariables/%s/" % (base_url, array_id, subvar_id)
-                    obj[key] = subvar_url
+                    new_obj[key] = subvar_url
                 else:
                     # Otherwise a regular variable references {"variable": alias}
                     var = variables.get(val)
@@ -716,10 +726,10 @@ def _process(obj, variables):
 
                     # TODO: We shouldn't stitch URLs together, use the API
                     if var.get('is_subvar'):
-                        obj[key] = '%svariables/%s/subvariables/%s/' \
+                        new_obj[key] = '%svariables/%s/subvariables/%s/' \
                                    % (base_url, var['parent_id'], var['id'])
                     else:
-                        obj[key] = '%svariables/%s/' % (base_url, var['id'])
+                        new_obj[key] = '%svariables/%s/' % (base_url, var['id'])
 
                     if var['type'] in ARRAY_TYPES:
                         subvariables = []
@@ -731,8 +741,19 @@ def _process(obj, variables):
                                 '%svariables/%s/subvariables/%s/'
                                 % (base_url, var['id'], subvar_id)
                             )
+            elif key == 'var':
+                var = variables.get(val)
+                if not var:
+                    raise ValueError("Invalid variable alias '%s'" % val)
+                if var.get('is_subvar'):
+                    parents_by_ids = {v['id']: v for v in variables.values() if not v.get('is_subvar')}
+                    parent = parents_by_ids[var['parent_id']]
+                    new_obj[key] = parent['alias']
+                    new_obj['axes'] = [val]
             elif key == 'function':
                 op = val
+        
+        obj = new_obj
 
         if subvariables:
             obj['subvariables'] = subvariables
@@ -778,11 +799,21 @@ def _process(obj, variables):
                     value['value'] = inner_value[0]
 
                 if len(subvariables) == 1:
+                    # obj['function'] = real_op
+                    # obj["args"] = [
+                    #     {'variable': subvariables[0]},
+                    #     value
+                    # ]
                     obj['function'] = real_op
-                    obj["args"] = [
-                        {'variable': subvariables[0]},
+                    obj['args'] = [
+                        {
+                            'var': var['alias'],
+                            'axes': [var['subreferences'][subvariables[0]]['alias']]
+                        },
                         value
                     ]
+                    # obj['args'][0]['axes'] = [var['subreferences'][subvariables[0]]['alias']]
+                    # obj['args'][1] = value
                 else:
                     obj = {
                         'function': expansion_op,
@@ -793,7 +824,7 @@ def _process(obj, variables):
                         [{
                             'function': real_op,
                             'args': [
-                                {'variable': subvar},
+                                {'var': var['alias'], 'axes': [var['subreferences'][subvar]['alias']]},
                                 value
                             ]
                         } for subvar in subvariables]

scrunch/tests/test_cubes.py

@@ -48,7 +48,7 @@ def test_pass_filter_expression(self, mock_fetch_cube):
             "function": ">",
             "args": [
                 {
-                    "variable": "https://test.crunch.io/api/datasets/123456/variables/0001/"
+                    "var": "var1_alias"
                 },
                 {"value": 1},
             ],