Merge pull request #372 from ecmwf-ifs/naml-ir-expr-visitor

IR: Move `expr_visitors` to `loki.ir`

loki/analyse/analyse_dataflow.py

@@ -10,10 +10,12 @@
 """
 
 from contextlib import contextmanager
-from loki.expression import FindVariables, Array, FindInlineCalls, ProcedureSymbol, FindTypedSymbols
+from loki.expression import Array, ProcedureSymbol
 from loki.tools import as_tuple, flatten
 from loki.types import BasicType
-from loki.ir import Visitor, Transformer
+from loki.ir import (
+    Visitor, Transformer, FindVariables, FindInlineCalls, FindTypedSymbols
+)
 from loki.subroutine import Subroutine
 from loki.tools.util import CaseInsensitiveDict
 

loki/analyse/util_polyhedron.py

@@ -7,13 +7,10 @@
 
 from typing import List
 import numpy as np
-from loki.ir import Loop
+from loki.ir import Loop, FindVariables
 from loki.expression import (
-    symbols as sym,
-    FindVariables,
-    simplify,
-    is_constant,
-    accumulate_polynomial_terms,
+    symbols as sym, simplify, is_constant,
+    accumulate_polynomial_terms
 )
 from loki.tools import as_tuple
 

loki/batch/item.py

@@ -11,10 +11,11 @@
 from loki.batch.configure import SchedulerConfig, ItemConfig
 from loki.frontend import REGEX, RegexParserClass
 from loki.expression import (
-    TypedSymbol, MetaSymbol, ProcedureSymbol, Variable, FindInlineCalls
+    TypedSymbol, MetaSymbol, ProcedureSymbol, Variable
 )
 from loki.ir import (
-    Import, CallStatement, TypeDef, ProcedureDeclaration, Interface, FindNodes
+    Import, CallStatement, TypeDef, ProcedureDeclaration, Interface,
+    FindNodes, FindInlineCalls
 )
 from loki.logging import warning
 from loki.module import Module

loki/batch/tests/test_scheduler.py

@@ -66,13 +66,11 @@
     ProcedureBindingItem, InterfaceItem, TypeDefItem, SFilter,
     ExternalItem, Transformation, Pipeline
 )
-from loki.expression import (
-    Scalar, Array, Literal, ProcedureSymbol, FindInlineCalls
-)
+from loki.expression import Scalar, Array, Literal, ProcedureSymbol
 from loki.frontend import (
     available_frontends, OMNI, OFP, FP, REGEX, HAVE_FP, HAVE_OFP, HAVE_OMNI
 )
-from loki.ir import nodes as ir, FindNodes
+from loki.ir import nodes as ir, FindNodes, FindInlineCalls
 from loki.transformations import (
     DependencyTransformation, ModuleWrapTransformation
 )

loki/expression/__init__.py

@@ -9,7 +9,6 @@
 <https://github.com/inducer/pymbolic>`_.
 """
 
-from loki.expression.expr_visitors import *  # noqa
 from loki.expression.symbols import *  # noqa
 from loki.expression.operations import *  # noqa
 from loki.expression.mappers import *  # noqa

loki/expression/parser.py

@@ -26,11 +26,9 @@
 except ImportError:
     FORTRAN_INTRINSIC_PROCEDURES = ()
 
-from loki.tools.util import CaseInsensitiveDict
-import loki.expression.symbols as sym
-import loki.expression.operations as sym_ops
-from loki.expression.expr_visitors import AttachScopes
+from loki.expression import symbols as sym, operations as sym_ops
 from loki.scope import Scope
+from loki.tools.util import CaseInsensitiveDict
 
 __all__ = ['ExpressionParser', 'parse_expr', 'FORTRAN_INTRINSIC_PROCEDURES', 'LokiEvaluationMapper']
 
@@ -519,6 +517,7 @@ def __call__(self, expr_str, scope=None, evaluate=False, strict=False, context=N
         :any:`Expression`
             The expression tree corresponding to the expression
         """
+        from loki.ir import AttachScopes  # pylint: disable=import-outside-toplevel,cyclic-import
         result = super().__call__(expr_str)
         context = context or {}
         context = CaseInsensitiveDict(context)

loki/expression/tests/test_expression.py

@@ -20,14 +20,14 @@
 )
 from loki.backend import cgen, fgen
 from loki.build import jit_compile, clean_test
-from loki.expression import (
-    symbols as sym, FindVariables, FindExpressions, FindInlineCalls,
-    SubstituteExpressions, AttachScopesMapper, parse_expr
-)
+from loki.expression import symbols as sym, parse_expr, AttachScopesMapper
 from loki.frontend import (
     available_frontends, OFP, OMNI, FP, HAVE_FP, parse_fparser_expression
 )
-from loki.ir import nodes as ir, FindNodes
+from loki.ir import (
+    nodes as ir, FindNodes, FindVariables, FindExpressions,
+    FindInlineCalls, SubstituteExpressions
+)
 from loki.tools import (
     filehash, stdchannel_redirected, stdchannel_is_captured
 )

loki/expression/tests/test_parser.py

@@ -11,10 +11,11 @@
 import pymbolic.mapper as pmbl_mapper
 
 from loki import Subroutine, Module, Scope
-from loki.expression import symbols as sym, FindVariables, parse_expr
+from loki.expression import symbols as sym, parse_expr
 from loki.frontend import (
     available_frontends, OMNI, HAVE_FP, parse_fparser_expression
 )
+from loki.ir import FindVariables
 
 
 # utility function to test parse_expr with different case

loki/frontend/fparser.py

@@ -28,13 +28,14 @@
 from loki import ir
 from loki.ir import (
     GenericVisitor, Transformer, FindNodes, attach_pragmas,
-    process_dimension_pragmas, detach_pragmas, pragmas_attached
+    process_dimension_pragmas, detach_pragmas, pragmas_attached,
+    AttachScopes
 )
 import loki.expression.symbols as sym
 from loki.expression.operations import (
     StringConcat, ParenthesisedAdd, ParenthesisedMul, ParenthesisedDiv, ParenthesisedPow
 )
-from loki.expression import ExpressionDimensionsMapper, AttachScopes, AttachScopesMapper
+from loki.expression import ExpressionDimensionsMapper, AttachScopesMapper
 from loki.logging import debug, perf, info, warning, error
 from loki.tools import (
     as_tuple, flatten, CaseInsensitiveDict, LazyNodeLookup, dict_override

loki/frontend/tests/test_omni.py

@@ -12,9 +12,8 @@
 import pytest
 
 from loki import Module, Subroutine
-from loki.expression import FindVariables
 from loki.frontend import OMNI, HAVE_OMNI
-from loki.ir import nodes as ir, FindNodes
+from loki.ir import nodes as ir, FindNodes, FindVariables
 
 
 @pytest.mark.skipif(not HAVE_OMNI, reason='Test tequires OMNI frontend.')

loki/ir/__init__.py

@@ -9,6 +9,7 @@
 The Loki internal representation (IR) and associated APIs for tree traversal.
 """
 
+from loki.ir.expr_visitors import *  # noqa
 from loki.ir.find import *  # noqa
 from loki.ir.ir_graph import *  # noqa
 from loki.ir.nodes import *  # noqa

loki/expression/tests/test_expr_visitors.py → loki/ir/tests/test_expr_visitors.py

@@ -8,12 +8,12 @@
 import pytest
 
 from loki import Sourcefile, Subroutine
-from loki.expression import (
-    symbols as sym, parse_expr, FindVariables, FindTypedSymbols,
+from loki.expression import symbols as sym, parse_expr
+from loki.frontend import available_frontends
+from loki.ir import (
+    nodes as ir, FindNodes, FindVariables, FindTypedSymbols,
     SubstituteExpressions, SubstituteStringExpressions
 )
-from loki.frontend import available_frontends
-from loki.ir import nodes as ir, FindNodes
 
 
 @pytest.mark.parametrize('frontend', available_frontends())

loki/ir/tests/test_visitor.py

@@ -16,12 +16,11 @@
 from loki.ir import (
     is_parent_of, is_child_of, FindNodes, FindScopes, Transformer,
     NestedTransformer, MaskedTransformer, NestedMaskedTransformer,
-    Stringifier
+    Stringifier, FindVariables, ExpressionFinder,
+    SubstituteExpressions
 )
 from loki.expression import (
-    symbols as sym, FindVariables, ExpressionFinder,
-    ExpressionCallbackMapper, ExpressionRetriever,
-    SubstituteExpressions
+    symbols as sym, ExpressionCallbackMapper, ExpressionRetriever
 )
 
 

loki/scope.py

@@ -296,7 +296,7 @@ def rescope_symbols(self):
         Make sure all symbols declared and used inside this node belong
         to a scope in the scope hierarchy
         """
-        from loki.expression import AttachScopes  # pylint: disable=import-outside-toplevel,cyclic-import
+        from loki.ir import AttachScopes  # pylint: disable=import-outside-toplevel,cyclic-import
         AttachScopes().visit(self, scope=self)
 
     def make_complete(self, **frontend_args):

loki/transformations/argument_shape.py

@@ -22,8 +22,10 @@
 
 
 from loki.batch import Transformation
-from loki.expression import Array, FindVariables, SubstituteExpressions
-from loki.ir import FindNodes, CallStatement, Transformer
+from loki.expression import Array
+from loki.ir import (
+    FindNodes, CallStatement, Transformer, FindVariables, SubstituteExpressions
+)
 from loki.tools import as_tuple, CaseInsensitiveDict
 from loki.types import BasicType
 

loki/transformations/array_indexing.py

@@ -16,12 +16,10 @@
 from loki.batch import Transformation, ProcedureItem
 from loki.logging import info
 from loki.analyse import dataflow_analysis_attached
-from loki.expression import (
-    symbols as sym, simplify, symbolic_op, FindVariables, SubstituteExpressions,
-    is_constant
-)
+from loki.expression import symbols as sym, simplify, symbolic_op, is_constant
 from loki.ir import (
-    Assignment, Loop, VariableDeclaration, FindNodes, Transformer, nodes as ir
+    nodes as ir, Assignment, Loop, VariableDeclaration, FindNodes,
+    Transformer, FindVariables, SubstituteExpressions
 )
 from loki.tools import as_tuple, CaseInsensitiveDict
 from loki.types import SymbolAttributes, BasicType

loki/transformations/block_index_transformations.py

@@ -7,17 +7,17 @@
 
 from loki.batch import Transformation, ProcedureItem
 from loki.ir import (
-        nodes as ir, FindNodes, Transformer, pragmas_attached,
-        pragma_regions_attached
+    nodes as ir, FindNodes, Transformer, pragmas_attached,
+    pragma_regions_attached, FindVariables, SubstituteExpressions,
+    AttachScopes
 )
+from loki.logging import warning
 from loki.module import Module
-from loki.tools import as_tuple # , CaseInsensitiveDict
+from loki.tools import as_tuple
 from loki.types import SymbolAttributes, BasicType
 from loki.expression import (
-        Variable, Array, RangeIndex, FindVariables, SubstituteExpressions,
-        symbols as sym, AttachScopes
+    symbols as sym, Variable, Array, RangeIndex
 )
-from loki.logging import warning
 from loki.transformations.sanitise import resolve_associates
 from loki.transformations.utilities import (
     recursive_expression_map_update, get_integer_variable,

loki/transformations/build_system/dependency.py

@@ -9,8 +9,9 @@
 
 from loki.backend import fgen
 from loki.batch import Transformation
-from loki.expression import FindInlineCalls
-from loki.ir import CallStatement, Import, Interface, FindNodes, Transformer
+from loki.ir import (
+    CallStatement, Import, Interface, FindNodes, FindInlineCalls, Transformer
+)
 from loki.logging import warning
 from loki.module import Module
 from loki.scope import Scope

loki/transformations/build_system/tests/test_dependency.py

@@ -9,9 +9,10 @@
 
 from loki import Sourcefile
 from loki.batch import Scheduler, SchedulerConfig
-from loki.expression import FindInlineCalls
 from loki.frontend import available_frontends, OMNI, OFP
-from loki.ir import FindNodes, CallStatement, Import, Interface, Intrinsic
+from loki.ir import (
+    FindNodes, CallStatement, Import, Interface, Intrinsic, FindInlineCalls
+)
 
 from loki.transformations import (
     DependencyTransformation, ModuleWrapTransformation

loki/transformations/data_offload.py

@@ -10,12 +10,11 @@
 
 from loki.analyse import dataflow_analysis_attached
 from loki.batch import Transformation, ProcedureItem, ModuleItem
-from loki.expression import (
-    Scalar, Array, FindInlineCalls, SubstituteExpressions
-)
+from loki.expression import Scalar, Array
 from loki.ir import (
     FindNodes, PragmaRegion, CallStatement, Pragma, Import, Comment,
-    Transformer, pragma_regions_attached, get_pragma_parameters
+    Transformer, pragma_regions_attached, get_pragma_parameters,
+    FindInlineCalls, SubstituteExpressions
 )
 from loki.logging import warning
 from loki.tools import as_tuple, flatten, CaseInsensitiveDict, CaseInsensitiveDefaultDict

loki/transformations/extract.py

@@ -6,11 +6,11 @@
 # nor does it submit to any jurisdiction.
 
 from loki.subroutine import Subroutine
-from loki.expression import (
-    FindVariables, FindInlineCalls, SubstituteExpressions,
-    DeferredTypeSymbol, Array
+from loki.expression import DeferredTypeSymbol, Array
+from loki.ir import (
+    CallStatement, Transformer, FindNodes, FindVariables,
+    FindInlineCalls, SubstituteExpressions
 )
-from loki.ir import CallStatement, Transformer, FindNodes
 from loki.types import DerivedType
 
 

loki/transformations/hoist_variables.py

@@ -83,13 +83,10 @@
 from collections import defaultdict
 
 from loki.batch import Transformation, ProcedureItem
-from loki.expression import (
-    symbols as sym, FindVariables, FindInlineCalls,
-    SubstituteExpressions, is_dimension_constant
-)
+from loki.expression import symbols as sym, is_dimension_constant
 from loki.ir import (
     CallStatement, Allocation, Deallocation, Transformer, FindNodes, Comment, Import,
-    Assignment
+    Assignment, FindVariables, FindInlineCalls, SubstituteExpressions
 )
 from loki.tools.util import is_iterable, as_tuple, CaseInsensitiveDict, flatten
 

loki/transformations/inline.py

@@ -16,12 +16,10 @@
 from loki.ir import (
     Import, Comment, Assignment, VariableDeclaration, CallStatement,
     Transformer, FindNodes, pragmas_attached, is_loki_pragma, Interface,
-    StatementFunction
-)
-from loki.expression import (
-    symbols as sym, FindVariables, FindInlineCalls, FindLiterals,
-    SubstituteExpressions, LokiIdentityMapper
+    StatementFunction, FindVariables, FindInlineCalls, FindLiterals,
+    SubstituteExpressions
 )
+from loki.expression import symbols as sym, LokiIdentityMapper
 from loki.types import BasicType
 from loki.tools import as_tuple, CaseInsensitiveDict
 from loki.logging import warning, error

loki/transformations/loop_blocking.py

@@ -6,10 +6,13 @@
 # nor does it submit to any jurisdiction.
 
 
-from loki.ir import nodes as ir, Transformer
+from loki.ir import (
+    nodes as ir, Transformer, FindVariables, SubstituteExpressions
+)
 from loki.subroutine import Subroutine
-from loki.expression import symbols as sym, parse_expr, FindVariables, \
-    SubstituteExpressions, ceil_division, iteration_index
+from loki.expression import (
+    symbols as sym, parse_expr, ceil_division, iteration_index
+)
 
 __all__ = ['split_loop', 'block_loop_arrays']
 

loki/transformations/pool_allocator.py

@@ -11,16 +11,16 @@
 from loki.batch import Transformation
 from loki.analyse import dataflow_analysis_attached
 from loki.expression import (
-    FindVariables, FindInlineCalls, Quotient, IntLiteral,
-    LogicLiteral, Variable, Array, Sum, Literal, Product, InlineCall,
-    Comparison, RangeIndex, Cast, ProcedureSymbol, LogicalNot,
-    simplify, SubstituteExpressions, is_dimension_constant,
+    Quotient, IntLiteral, LogicLiteral, Variable, Array, Sum, Literal,
+    Product, InlineCall, Comparison, RangeIndex, Cast,
+    ProcedureSymbol, LogicalNot, simplify, is_dimension_constant,
     DetachScopesMapper
 )
 from loki.ir import (
-    FindNodes, Transformer, Intrinsic, Assignment, Conditional,
-    CallStatement, Import, Allocation, Deallocation, Loop, Pragma,
-    Interface, get_pragma_parameters
+    FindNodes, FindVariables, FindInlineCalls, Transformer, Intrinsic,
+    Assignment, Conditional, CallStatement, Import, Allocation,
+    Deallocation, Loop, Pragma, Interface, get_pragma_parameters,
+    SubstituteExpressions
 )
 from loki.logging import warning, debug
 from loki.tools import as_tuple

loki/transformations/raw_stack_allocator.py

@@ -17,7 +17,7 @@
 )
 from loki.expression.symbolic import is_dimension_constant, simplify
 from loki.expression.mappers import DetachScopesMapper
-from loki.expression.expr_visitors import FindVariables, SubstituteExpressions
+from loki.ir.expr_visitors import FindVariables, SubstituteExpressions
 from loki.ir.nodes import Assignment, CallStatement, Pragma
 from loki.ir.find import FindNodes
 from loki.ir.transformer import Transformer