From 9f9ca78a46cbc297f227161074657df60feb413d Mon Sep 17 00:00:00 2001
From: Michael Lange <Michael.Lange@ecmwf.int>
Date: Tue, 7 May 2024 12:21:50 +0000
Subject: [PATCH] SCCRawStack: Re-organise import and make static helper a
 property

---
 loki/transformations/raw_stack_allocator.py   | 42 +++++++++++--------
 .../tests/test_raw_stack_allocator.py         | 18 ++++----
 2 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/loki/transformations/raw_stack_allocator.py b/loki/transformations/raw_stack_allocator.py
index ccadadb2e..8b7983340 100644
--- a/loki/transformations/raw_stack_allocator.py
+++ b/loki/transformations/raw_stack_allocator.py
@@ -7,21 +7,23 @@
 
 import re
 
-from loki.batch import Transformation
-from loki.expression import Array, Scalar
-from loki.types import BasicType
 from loki.analyse import dataflow_analysis_attached
-from loki.types import SymbolAttributes
-from loki.ir import Assignment, CallStatement, Pragma
-from loki.tools import as_tuple
-from loki.ir import FindNodes, Transformer
-from loki.batch import ProcedureItem
+from loki.backend.fgen import fgen
+from loki.batch.item import ProcedureItem
+from loki.batch.transformation import Transformation
+from loki.expression.symbols import (
+    Array, Scalar, Variable, Literal, Product, Sum, InlineCall,
+    IntLiteral, RangeIndex, DeferredTypeSymbol
+)
 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.backend.fgen import fgen
-from loki.expression.symbols import (
-    Variable, Literal, Product, Sum, InlineCall, IntLiteral, RangeIndex, DeferredTypeSymbol)
+from loki.ir.nodes import Assignment, CallStatement, Pragma
+from loki.ir.find import FindNodes
+from loki.ir.transformer import Transformer
+from loki.tools import as_tuple
+from loki.types import BasicType, SymbolAttributes
+
 
 __all__ = ['TemporariesRawStackTransformation']
 
@@ -74,6 +76,12 @@ class TemporariesRawStackTransformation(Transformation):
     # Traverse call tree in reverse when using Scheduler
     reverse_traversal = True
 
+    type_name_dict = {
+        BasicType.REAL: {'kernel': 'P', 'driver': 'Z'},
+        BasicType.LOGICAL: {'kernel': 'LD', 'driver': 'LL'},
+        BasicType.INTEGER: {'kernel': 'K', 'driver': 'I'}
+    }
+
     def __init__(self, block_dim, horizontal,
                  stack_name='STACK',
                  local_int_var_name_pattern='JD_{name}',
@@ -88,13 +96,11 @@ def __init__(self, block_dim, horizontal,
         if key:
             self._key = key
 
-
-    int_type = SymbolAttributes(dtype=BasicType.INTEGER, kind=DeferredTypeSymbol('JPIM'))
-
-    type_name_dict = {BasicType.REAL: {'kernel': 'P', 'driver': 'Z'},
-                      BasicType.LOGICAL: {'kernel': 'LD', 'driver': 'LL'},
-                      BasicType.INTEGER: {'kernel': 'K', 'driver': 'I'}}
-
+    @property
+    def int_type(self):
+        return SymbolAttributes(
+            dtype=BasicType.INTEGER, kind=DeferredTypeSymbol('JPIM')
+        )
 
     def transform_subroutine(self, routine, **kwargs):
 
diff --git a/loki/transformations/tests/test_raw_stack_allocator.py b/loki/transformations/tests/test_raw_stack_allocator.py
index cb3fda566..c4ea4c455 100644
--- a/loki/transformations/tests/test_raw_stack_allocator.py
+++ b/loki/transformations/tests/test_raw_stack_allocator.py
@@ -9,20 +9,18 @@
 
 import pytest
 
-from loki.tools import gettempdir
-from loki.dimension import Dimension
+from loki.backend import fgen
 from loki.batch import Scheduler, SchedulerConfig, ProcedureItem
-from loki.frontend.util import OMNI
-from loki.backend.fgen import fgen
-from loki.types import BasicType
+from loki.dimension import Dimension
+from loki.expression import DeferredTypeSymbol, InlineCall, IntLiteral
+from loki.frontend import available_frontends, OMNI
 from loki.ir import FindNodes, CallStatement, Assignment, Pragma
 from loki.sourcefile import Sourcefile
-from loki.expression.symbols import DeferredTypeSymbol, InlineCall, IntLiteral
-from loki.transform.transform_array_indexing import normalize_range_indexing
-
-from conftest import available_frontends
+from loki.tools import gettempdir
+from loki.types import BasicType
 
-from transformations.raw_stack_allocator import TemporariesRawStackTransformation
+from loki.transformations.array_indexing import normalize_range_indexing
+from loki.transformations.raw_stack_allocator import TemporariesRawStackTransformation
 
 
 @pytest.fixture(scope='module', name='block_dim')