parse_expr: allow for nested derived types

Author: MichaelSt98

loki/expression/parser.py

@@ -122,7 +122,7 @@ def map_slice(self, expr, *args, **kwargs):
 
     def map_variable(self, expr, *args, **kwargs):
         parent = kwargs.pop('parent', None)
-        return sym.Variable(name=expr.name, parent=parent) # , **kwargs)
+        return sym.Variable(name=expr.name, parent=parent)
 
     def map_algebraic_leaf(self, expr, *args, **kwargs):
         if str(expr).isnumeric():
@@ -255,30 +255,39 @@ def map_call_with_kwargs(self, expr):
         return self.rec(expr.function)(*args, **kwargs)
 
     def map_lookup(self, expr):
-        try:
-            if isinstance(expr.name, pmbl.Variable):
-                name = expr.name.name
-                return self.case_insensitive_getattr(self.rec(expr.aggregate), name)
-            if isinstance(expr.name, pmbl.Call):
-                name = expr.name.function.name
-                if callable(self.case_insensitive_getattr(self.rec(expr.aggregate), name)):
-                    return self.case_insensitive_getattr(self.rec(expr.aggregate),
-                            name)(*[self.rec(par) for par in expr.name.parameters])
-                return self._evaluate_array(self.case_insensitive_getattr(self.rec(expr.aggregate), name),
-                        [self.rec(par) for par in expr.name.parameters])
-            if isinstance(expr.name, pmbl.CallWithKwargs):
-                name = expr.name.function.name
-                args = [self.rec(par) for par in expr.name.parameters]
-                kwargs = {
-                    k: self.rec(v)
-                    for k, v in expr.name.kw_parameters.items()}
-                kwargs = CaseInsensitiveDict(kwargs)
-                return self.case_insensitive_getattr(self.rec(expr.aggregate), name)(*args, **kwargs)
 
+        def rec_lookup(expr, obj, name):
+            return expr.name, self.case_insensitive_getattr(obj, name)
+
+        try:
+            current_expr = expr
+            obj = self.rec(expr.aggregate)
+            while isinstance(current_expr.name, pmbl.Lookup):
+                current_expr, obj = rec_lookup(current_expr, obj, current_expr.name.aggregate.name)
+            if isinstance(current_expr.name, pmbl.Variable):
+                _, obj = rec_lookup(current_expr, obj, current_expr.name.name)
+                return obj
+            if isinstance(current_expr.name, pmbl.Call):
+                name = current_expr.name.function.name
+                _, obj = rec_lookup(current_expr, obj, name)
+                if callable(obj):
+                    return obj(*[self.rec(par) for par in current_expr.name.parameters])
+                return self._evaluate_array(obj, [self.rec(par) for par in current_expr.name.parameters])
+            if isinstance(current_expr.name, pmbl.CallWithKwargs):
+                name = current_expr.name.function.name
+                _, obj = rec_lookup(current_expr, obj, name)
+                args = [self.rec(par) for par in current_expr.name.parameters]
+                kwargs = CaseInsensitiveDict(
+                    (k, self.rec(v))
+                    for k, v in current_expr.name.kw_parameters.items()
+                )
+                return obj(*args, **kwargs)
         except Exception as e:
             if self.strict:
                 raise e
             return expr
+        if self.strict:
+            raise NotImplementedError
         return expr
 
 

loki/expression/tests/test_expression.py

@@ -13,6 +13,7 @@
 import numpy as np
 
 import pymbolic.primitives as pmbl
+import pymbolic.mapper as pmbl_mapper
 
 from loki import (
     Sourcefile, Subroutine, Module, Scope, BasicType,
@@ -2109,7 +2110,7 @@ def test_expression_parser_evaluate(case):
     context = {'a': 6}
     test_str = '1 + 1 + a + some_func(a, 10)'
     parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
-    with pytest.raises(Exception):
+    with pytest.raises(pmbl_mapper.evaluator.UnknownVariableError):
         parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, strict=True, context=context)
     parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, strict=False, context=context)
     assert str(parsed).lower().replace(' ', '') == '8+some_func(6,10)'
@@ -2134,7 +2135,23 @@ def some_func(a, b, c=None):
     parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
     assert str(parsed).lower().replace(' ', '') == '(13+c+1)/(c+1)'
 
+    class BarBarBar:
+        val_barbarbar = 5
+
+    class BarBar:
+        barbarbar = BarBarBar()
+        val_barbar = -3
+        def barbar_func(self, a):
+            return a - 1
+
+    class Bar:
+        barbar = BarBar()
+        val_bar = 5
+        def bar_func(self, a):
+            return a**2
+
     class Foo:
+        bar = Bar() # pylint: disable=disallowed-name
         val3 = 1
         arr = [[1, 2], [3, 4]]
         def __init__(self, _val1, _val2):
@@ -2150,18 +2167,33 @@ def static_func(a):
     test_str = 'foo%val1 + foo%val2 + foo%val3'
     parsed = parse_expr(convert_to_case(f'{test_str}', mode=case))
     assert str(parsed).lower().replace(' ', '') == 'foo%val1+foo%val2+foo%val3'
-    with pytest.raises(Exception):
+    with pytest.raises(pmbl_mapper.evaluator.UnknownVariableError):
         parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, strict=True)
     parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
     assert parsed == 6
     test_str = 'foo%val1 + foo%some_func(1, 2) + foo%static_func_2(3)'
     parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
     assert str(parsed).lower().replace(' ', '') == '5+foo%static_func_2(3)'
-    with pytest.raises(Exception):
+    with pytest.raises(pmbl_mapper.evaluator.UnknownVariableError):
         parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, strict=True)
     test_str = 'foo%val1 + foo%some_func(1, 2) + foo%static_func(3) + foo%arr(1, 2)'
     parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context, strict=True)
     assert parsed == 13
     test_str = 'foo%val1 + foo%some_func(1, b=2) + foo%static_func(a=3) + foo%arr(1, 2)'
     parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context, strict=True)
     assert parsed == 13
+    test_str = 'foo%bar%val_bar + 1'
+    parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
+    assert parsed == 6
+    test_str = 'foo%bar%bar_func(2) + 1'
+    parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
+    assert parsed == 5
+    test_str = 'foo%bar%barbar%val_barbar + 1'
+    parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
+    assert parsed == -2
+    test_str = 'foo%bar%barbar%barbar_func(0) + 1'
+    parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
+    assert parsed == 0
+    test_str = 'foo%bar%barbar%barbarbar%val_barbarbar + 1'
+    parsed = parse_expr(convert_to_case(f'{test_str}', mode=case), evaluate=True, context=context)
+    assert parsed == 6