Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix #120: Test and fix constructor arg conflict in FindExpressionRoot #121

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion loki/expression/expr_visitors.py
Original file line number Diff line number Diff line change
Expand Up @@ -256,7 +256,10 @@ class FindExpressionRoot(ExpressionFinder):
"""
def __init__(self, expr, recurse_to_parent=True, **kwargs):
self._retriever = ExpressionRetriever(lambda e: e is expr, recurse_to_parent=recurse_to_parent)
super().__init__(unique=False, retrieve=lambda e: e if self._retriever.retrieve(e) else (), **kwargs)
if kwargs.get('unique'):
raise ValueError('FindExpressionRoot requires unique=False')
kwargs['unique'] = False
super().__init__(retrieve=lambda e: e if self._retriever.retrieve(e) else (), **kwargs)


class SubstituteExpressions(Transformer):
Expand Down
24 changes: 24 additions & 0 deletions tests/test_visitor.py
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,30 @@ def test_find_expression_root(frontend):
assert cast_root[0].source.lines == (13, 13)


@pytest.mark.parametrize('frontend', available_frontends())
def test_find_expression_root_constructor_args(frontend):
"""
Test correct handling for various constructor arguments
"""
fcode = """
subroutine my_routine
implicit none
integer :: i
i = 1 + 1
end subroutine my_routine
""".strip()

routine = Subroutine.from_source(fcode, frontend=frontend)
exprs = FindExpressions().visit(routine.body)
some_expr = [expr for expr in exprs if isinstance(expr, IntLiteral)][0]

with pytest.raises(ValueError):
FindExpressionRoot(some_expr, unique=True).visit(routine.body)

expr_root = FindExpressionRoot(some_expr, unique=False).visit(routine.body)
assert expr_root == (routine.body.body[0].rhs,)


@pytest.mark.parametrize('frontend', available_frontends())
def test_find_variables_associates(frontend):
"""
Expand Down