bug: added AcousticWaveSolverType

mrava87 · Aug 13, 2024 · 1b3ad18 · 1b3ad18
1 parent 0c2e852
commit 1b3ad18
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/pylops/waveeqprocessing/twoway.py b/pylops/waveeqprocessing/twoway.py
@@ -1,6 +1,6 @@
 __all__ = ["AcousticWave2D"]
 
-from typing import Tuple
+from typing import Any, NewType, Tuple
 
 import numpy as np
 
@@ -15,6 +15,10 @@
     from examples.seismic import AcquisitionGeometry, Model
     from examples.seismic.acoustic import AcousticWaveSolver
     from examples.seismic.utils import _CustomSource
+else:
+    AcousticWaveSolver = Any
+
+AcousticWaveSolverType = NewType("AcousticWaveSolver", AcousticWaveSolver)
 
 
 class AcousticWave2D(LinearOperator):
@@ -285,7 +289,7 @@ def srcillumination_allshots(self, savewav: bool = False) -> None:
                 self.src_wavefield.append(src_wav)
             self.src_illumination += src_ill
 
-    def _born_oneshot(self, solver: AcousticWaveSolver, dm: NDArray) -> NDArray:
+    def _born_oneshot(self, solver: AcousticWaveSolverType, dm: NDArray) -> NDArray:
         """Born modelling for one shot
 
         Parameters