Fixed typing

src/parakeet/_run.py

@@ -14,10 +14,9 @@
 import parakeet.metadata
 import parakeet.sample
 import parakeet.simulate
+from parakeet.config import Device
 from functools import singledispatch
 
-Device = parakeet.config.Device
-
 
 __all__ = ["run"]
 
@@ -33,7 +32,7 @@ def run(
     exit_wave_file: str,
     optics_file: str,
     image_file: str,
-    device: str = None,
+    device: Device = None,
     nproc: int = None,
     gpu_id: list = None,
     steps: list = None,

src/parakeet/analyse/_correct.py

@@ -13,6 +13,7 @@
 import random
 import parakeet.microscope
 import parakeet.sample
+from parakeet.config import Device
 from functools import singledispatch
 
 
@@ -29,7 +30,7 @@ def correct(
     image_file: str,
     corrected_file: str,
     num_defocus: int = 1,
-    device: str = None,
+    device: Device = None,
 ):
     """
     Correct the images using 3D CTF correction

src/parakeet/analyse/_reconstruct.py

@@ -13,6 +13,7 @@
 import random
 import parakeet.microscope
 import parakeet.sample
+from parakeet.config import Device
 from functools import singledispatch
 
 
@@ -24,7 +25,7 @@
 
 
 @singledispatch
-def reconstruct(config_file, image_file: str, rec_file: str, device: str = None):
+def reconstruct(config_file, image_file: str, rec_file: str, device: Device = None):
     """
     Reconstruct the volume
 

src/parakeet/simulate/_cbed.py

@@ -21,6 +21,7 @@
 import parakeet.io
 import parakeet.sample
 import parakeet.simulate
+from parakeet.config import Device
 from parakeet.simulate.simulation import Simulation
 from parakeet.microscope import Microscope
 from parakeet.scan import Scan
@@ -33,9 +34,6 @@
 __all__ = ["cbed"]
 
 
-Device = parakeet.config.Device
-Sample = parakeet.sample.Sample
-
 # Get the logger
 logger = logging.getLogger(__name__)
 
@@ -412,7 +410,7 @@ def cbed(
     config_file,
     sample_file: str,
     image_file: str,
-    device: str = None,
+    device: Device = None,
     nproc: int = None,
     gpu_id: list = None,
 ):

src/parakeet/simulate/_exit_wave.py

@@ -21,6 +21,7 @@
 import parakeet.io
 import parakeet.sample
 import parakeet.simulate
+from parakeet.config import Device
 from parakeet.simulate.simulation import Simulation
 from parakeet.microscope import Microscope
 from functools import singledispatch
@@ -414,7 +415,7 @@ def exit_wave(
     config_file,
     sample_file: str,
     exit_wave_file: str,
-    device: str = None,
+    device: Device = None,
     nproc: int = None,
     gpu_id: list = None,
 ):

src/parakeet/simulate/_optics.py

@@ -20,6 +20,7 @@
 import parakeet.inelastic
 import parakeet.io
 import parakeet.sample
+from parakeet.config import Device
 from parakeet.microscope import Microscope
 from parakeet.scan import Scan
 from functools import singledispatch
@@ -525,7 +526,7 @@ def optics(
     config_file,
     exit_wave_file: str,
     optics_file: str,
-    device: str = None,
+    device: Device = None,
     nproc: int = None,
     gpu_id: list = None,
 ):

src/parakeet/simulate/_potential.py

@@ -19,6 +19,7 @@
 import parakeet.futures
 import parakeet.inelastic
 import parakeet.sample
+from parakeet.config import Device
 from parakeet.sample import Sample
 from parakeet.microscope import Microscope
 from parakeet.scan import Scan
@@ -244,7 +245,7 @@ def potential(
     config_file,
     sample_file: str,
     potential_prefix: str,
-    device: str = None,
+    device: Device = None,
     nproc: int = None,
     gpu_id: list = None,
 ):