From 96aec6e5faf227901159c756f57e99afa2d7e312 Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Fri, 14 Jan 2022 17:39:54 +0100
Subject: [PATCH] Improve typing for protocol class

---
 kasa/protocol.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/kasa/protocol.py b/kasa/protocol.py
index 54fea0803..14de9c6c6 100755
--- a/kasa/protocol.py
+++ b/kasa/protocol.py
@@ -15,7 +15,7 @@
 import logging
 import struct
 from pprint import pformat as pf
-from typing import Dict, Optional, Union
+from typing import Dict, Generator, Optional, Union
 
 from .exceptions import SmartDeviceException
 
@@ -107,7 +107,7 @@ async def _execute_query(self, request: str) -> Dict:
 
         return json_payload
 
-    async def close(self):
+    async def close(self) -> None:
         """Close the connection."""
         writer = self.writer
         self.reader = self.writer = None
@@ -116,7 +116,7 @@ async def close(self):
             with contextlib.suppress(Exception):
                 await writer.wait_closed()
 
-    def _reset(self):
+    def _reset(self) -> None:
         """Clear any varibles that should not survive between loops."""
         self.reader = self.writer = self.loop = self.query_lock = None
 
@@ -154,13 +154,13 @@ async def _query(self, request: str, retry_count: int, timeout: int) -> Dict:
         await self.close()
         raise SmartDeviceException("Query reached somehow to unreachable")
 
-    def __del__(self):
+    def __del__(self) -> None:
         if self.writer and self.loop and self.loop.is_running():
             self.writer.close()
         self._reset()
 
     @staticmethod
-    def _xor_payload(unencrypted):
+    def _xor_payload(unencrypted: bytes) -> Generator[int, None, None]:
         key = TPLinkSmartHomeProtocol.INITIALIZATION_VECTOR
         for unencryptedbyte in unencrypted:
             key = key ^ unencryptedbyte
@@ -179,7 +179,7 @@ def encrypt(request: str) -> bytes:
         )
 
     @staticmethod
-    def _xor_encrypted_payload(ciphertext):
+    def _xor_encrypted_payload(ciphertext: bytes) -> Generator[int, None, None]:
         key = TPLinkSmartHomeProtocol.INITIALIZATION_VECTOR
         for cipherbyte in ciphertext:
             plainbyte = key ^ cipherbyte