misc: Various fixes

Author: JDBetteridge

examples/seismic/tti/tti_example.py

@@ -1,5 +1,6 @@
+from contextlib import suppress
+
 import numpy as np
-from contrextlib import suppress
 
 with suppress(ImportError):
     import pytest

examples/seismic/tutorials/07.1_dispersion_relation.ipynb

@@ -1132,10 +1132,10 @@
     "x = np.linspace(0, np.pi, 201)\n",
     "m = np.arange(1, len(fornberg) + 1)\n",
     "y_fornberg = - fornberg[0] - 2*np.sum(\n",
-    "    [a_ * np.cos(m_*x) for a_, m_ in zip(fornberg[1:], m, strict=True)], axis=0\n",
+    "    [a_ * np.cos(m_*x) for a_, m_ in zip(fornberg[1:], m, strict=False)], axis=0\n",
     ")\n",
     "y_drp1 = - drp_stencil1[0] - 2*np.sum(\n",
-    "    [a_ * np.cos(m_*x) for a_, m_ in zip(drp_stencil1[1:], m, strict=True)], axis=0\n",
+    "    [a_ * np.cos(m_*x) for a_, m_ in zip(drp_stencil1[1:], m, strict=False)], axis=0\n",
     ")\n",
     "\n",
     "fig, ax = plt.subplots(1, 2)\n",

tests/test_caching.py

@@ -678,7 +678,8 @@ def test_sparse_function(self, operate_on_empty_cache):
         ncreated = 2+1+2+2+2+1+4
         # Note that injection is now lazy so no new symbols should be created
         assert len(_SymbolCache) == cur_cache_size
-        _ = i.evaluate
+        # The expression is not redundant, but storing it changes the symbol count
+        i.evaluate  # noqa: B018
 
         assert len(_SymbolCache) == cur_cache_size + ncreated
 

tests/test_dimension.py

@@ -1492,8 +1492,29 @@ def test_stepping_dim_in_condition_lowering(self):
         op.apply(time_M=threshold+3)
         assert np.all(g.data[0, :, :] == threshold)
         assert np.all(g.data[1, :, :] == threshold + 1)
-        assert 'if (g[t0][x + 1][y + 1] <= 10)\n' + \
-            '{\n g[t1][x + 1][y + 1] = g[t0][x + 1][y + 1] + 1' in str(op.ccode)
+
+        # We want to assert that the following snippet:
+        # ```
+        #   if (g[t0][x + 1][y + 1] <= 10)
+        #   {
+        #     g[t1][x + 1][y + 1] = g[t0][x + 1][y + 1] + 1
+        # ```
+        # is in the generated code, but indentation etc. seems to vary.
+        part1 = 'if (g[t0][x + 1][y + 1] <= 10)\n'
+        part2 = 'g[t1][x + 1][y + 1] = g[t0][x + 1][y + 1] + 1'
+        whole_code = str(op.ccode)
+
+        try:
+            loc = whole_code.find(part1)
+            assert loc != -1
+            assert whole_code.find(part2, loc + len(part1)) != -1
+        except AssertionError:
+            # Try the alternative string
+            part1 = 'if (gL0(t0, x + 1, y + 1) <= 10)\n'
+            part2 = 'gL0(t1, x + 1, y + 1) = gL0(t0, x + 1, y + 1) + 1'
+            loc = whole_code.find(part1)
+            assert loc != -1
+            assert whole_code.find(part2, loc + len(part1)) != -1
 
     def test_expr_like_lowering(self):
         """

tests/test_dle.py

@@ -21,8 +21,10 @@
 
 
 def get_blocksizes(op, opt, grid, blockshape, level=0):
-    blocksizes = {f'{d}0_blk{level}_size': v
-                  for d, v in zip(grid.dimensions, blockshape, strict=True)}
+    blocksizes = {
+        f'{d}0_blk{level}_size': v
+        for d, v in zip(grid.dimensions, blockshape, strict=False)
+    }
     blocksizes = {k: v for k, v in blocksizes.items() if k in op._known_arguments}
     # Sanity check
     if grid.dim == 1 or len(blockshape) == 0: