Fix lint

kroo · May 13, 2021 · 05bc495 · 05bc495
1 parent 1411c2b
commit 05bc495
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/tests/test_wyzecam_doorbell.py b/tests/test_wyzecam_doorbell.py
@@ -1,12 +1,12 @@
 import pytest
-import wyzecam
-from wyzecam import WyzeAccount, WyzeCamera, WyzeIOTC, WyzeIOTCSessionState
+from wyzecam.api_models import WyzeAccount, WyzeCamera
+from wyzecam.iotc import WyzeIOTC, WyzeIOTCSessionState
 
 
 @pytest.mark.usefixtures("iotc", "account", "camera")
 def test_connect_and_auth(
     iotc: WyzeIOTC, account: WyzeAccount, camera: WyzeCamera
-):
+) -> None:
     camera.product_model = "WYZEDB3"
     with iotc.connect_and_auth(account, camera) as session:
         assert session.state == WyzeIOTCSessionState.AUTHENTICATION_SUCCEEDED

diff --git a/wyzecam/iotc.py b/wyzecam/iotc.py
@@ -62,7 +62,7 @@ def __init__(
         tutk_platform_lib: Optional[Union[str, CDLL]] = None,
         udp_port: Optional[int] = None,
         max_num_av_channels: Optional[int] = None,
-        debug=False,
+        debug: bool = False,
     ) -> None:
         """Construct a WyzeIOTC session object
 
@@ -310,7 +310,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     def recv_video_data(
         self,
-    ) -> Iterator[Tuple[Optional[bytes], tutk.FrameInfoStruct]]:
+    ) -> Iterator[
+        Tuple[
+            Optional[bytes], Union[tutk.FrameInfoStruct, tutk.FrameInfo3Struct]
+        ]
+    ]:
         """A generator for returning raw video frames!
 
         By iterating over the return value of this function, you will
@@ -380,7 +384,11 @@ def recv_video_data(
 
     def recv_video_frame(
         self,
-    ) -> Iterator[Tuple["av.VideoFrame", tutk.FrameInfoStruct]]:
+    ) -> Iterator[
+        Tuple[
+            "av.VideoFrame", Union[tutk.FrameInfoStruct, tutk.FrameInfo3Struct]
+        ]
+    ]:
         """A generator for returning decoded video frames!
 
         By iterating over the return value of this function, you will conveniently
@@ -424,7 +432,9 @@ def recv_video_frame(
 
     def recv_video_frame_ndarray(
         self,
-    ) -> Iterator[Tuple["np.ndarray", tutk.FrameInfoStruct]]:
+    ) -> Iterator[
+        Tuple["np.ndarray", Union[tutk.FrameInfoStruct, tutk.FrameInfo3Struct]]
+    ]:
         """A generator for returning decoded video frames!
 
         By iterating over the return value of this function, you will conveniently
@@ -470,7 +480,13 @@ def recv_video_frame_ndarray_with_stats(
         draw_stats: Optional[
             str
         ] = "{width}x{height} {kilobytes_per_second} kB/s {frames_per_second} FPS",
-    ) -> Iterator[Tuple["np.ndarray", tutk.FrameInfoStruct, Dict[str, int]]]:
+    ) -> Iterator[
+        Tuple[
+            "np.ndarray",
+            Union[tutk.FrameInfoStruct, tutk.FrameInfo3Struct],
+            Dict[str, int],
+        ]
+    ]:
         """
         Does everything recv_video_frame_ndarray does, but also computes a number
         of useful / interesting debug metrics including effective framerate, bitrate,