Add type hints

instamatic-dev · Oct 15, 2024 · 192d1d7 · 192d1d7
1 parent 240d336
commit 192d1d7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/instamatic/TEMController/TEMController.py b/src/instamatic/TEMController/TEMController.py
@@ -11,6 +11,7 @@
 from instamatic.formats import write_tiff
 from instamatic.image_utils import rotate_image
 from instamatic.TEMController.microscope_base import MicroscopeBase
+from instamatic.camera.camera_base import CameraBase
 
 from .deflectors import *
 from .lenses import *
@@ -90,7 +91,7 @@ class TEMController:
     cam: Camera control object (see instamatic.camera) [optional]
     """
 
-    def __init__(self, tem: MicroscopeBase, cam=None):
+    def __init__(self, tem: MicroscopeBase, cam: CameraBase = None):
         super().__init__()
 
         self._executor = ThreadPoolExecutor(max_workers=1)

diff --git a/src/instamatic/camera/videostream.py b/src/instamatic/camera/videostream.py
@@ -2,6 +2,7 @@
 import threading
 
 from .camera import Camera
+from instamatic.camera.camera_base import CameraBase
 
 
 class ImageGrabber:
@@ -15,7 +16,7 @@ class ImageGrabber:
     routine.
     """
 
-    def __init__(self, cam, callback, frametime: float = 0.05):
+    def __init__(self, cam: CameraBase, callback, frametime: float = 0.05):
         super().__init__()
 
         self.callback = callback
@@ -117,7 +118,7 @@ def send_frame(self, frame, acquire=False):
             self.frame = frame
             self.grabber.lock.release()
 
-    def setup_grabber(self):
+    def setup_grabber(self) -> ImageGrabber:
         grabber = ImageGrabber(self.cam, callback=self.send_frame, frametime=self.frametime)
         atexit.register(grabber.stop)
         return grabber