From 0bd0312ba47883946cdff0e917b176ee040493d0 Mon Sep 17 00:00:00 2001
From: James Parkhurst <james.parkhurst@diamond.ac.uk>
Date: Fri, 20 Oct 2023 09:39:04 +0100
Subject: [PATCH] Fixed typing

---
 src/parakeet/_run.py                 | 5 ++---
 src/parakeet/analyse/_correct.py     | 3 ++-
 src/parakeet/analyse/_reconstruct.py | 3 ++-
 src/parakeet/simulate/_cbed.py       | 6 ++----
 src/parakeet/simulate/_exit_wave.py  | 3 ++-
 src/parakeet/simulate/_optics.py     | 3 ++-
 src/parakeet/simulate/_potential.py  | 3 ++-
 7 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/parakeet/_run.py b/src/parakeet/_run.py
index 38d2ccc..447b2b4 100644
--- a/src/parakeet/_run.py
+++ b/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,
diff --git a/src/parakeet/analyse/_correct.py b/src/parakeet/analyse/_correct.py
index e75e788..3b12142 100644
--- a/src/parakeet/analyse/_correct.py
+++ b/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
diff --git a/src/parakeet/analyse/_reconstruct.py b/src/parakeet/analyse/_reconstruct.py
index 8154a6a..eec9612 100644
--- a/src/parakeet/analyse/_reconstruct.py
+++ b/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
 
diff --git a/src/parakeet/simulate/_cbed.py b/src/parakeet/simulate/_cbed.py
index 05f1549..eb369ca 100644
--- a/src/parakeet/simulate/_cbed.py
+++ b/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,
 ):
diff --git a/src/parakeet/simulate/_exit_wave.py b/src/parakeet/simulate/_exit_wave.py
index a7c4653..d381fe0 100644
--- a/src/parakeet/simulate/_exit_wave.py
+++ b/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,
 ):
diff --git a/src/parakeet/simulate/_optics.py b/src/parakeet/simulate/_optics.py
index b8ce4f3..0e684a7 100644
--- a/src/parakeet/simulate/_optics.py
+++ b/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,
 ):
diff --git a/src/parakeet/simulate/_potential.py b/src/parakeet/simulate/_potential.py
index 32a82a4..07757e3 100644
--- a/src/parakeet/simulate/_potential.py
+++ b/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,
 ):