fix import

adrn · Nov 27, 2019 · b4c216c · b4c216c
1 parent 6869894
commit b4c216c
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/schwimmbad/mpi.py b/schwimmbad/mpi.py
@@ -8,13 +8,17 @@
 # Still make it a global to avoid messing up other things.
 MPI = None
 
-# Filled in import_mpi below
-MPIPoolExecutor = object
-
 # Project
 from . import log, _VERBOSE
 from .pool import BasePool
 
+MPIPoolExecutor = object
+try:
+    from mpi4py.futures import MPIPoolExecutor as tmp
+    MPIPoolExecutor = tmp
+except ImportError:
+    pass
+
 __all__ = ['MPIPool', 'MPIAsyncPool']
 
 
@@ -32,9 +36,6 @@ def _import_mpi(quiet=False, use_dill=False):
             _MPI.pickle.__init__(dill.dumps, dill.loads, dill.HIGHEST_PROTOCOL)
         MPI = _MPI
 
-        from mpi4py.futures import MPIPoolExecutor as tmp
-        MPIPoolExecutor = tmp
-
     except ImportError:
         if not quiet:
             # Re-raise with a more user-friendly error: