Type hinting for _logger.py

croningp · May 11, 2020 · 10c435c · 10c435c
1 parent 37b0b4b
commit 10c435c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/pycont/_logger.py b/pycont/_logger.py
@@ -3,5 +3,5 @@
 __logger_root_name__ = 'pycont'
 
 
-def create_logger(name):
+def create_logger(name: str) -> logging.Logger:
     return logging.getLogger(__logger_root_name__ + '.' + name)
diff --git a/pycont/pump_protocol.py b/pycont/pump_protocol.py
@@ -7,6 +7,8 @@
 
 """
 # -*- coding: utf-8 -*-
+from typing import List, Union
+
 from ._logger import create_logger
 
 from . import dtprotocol
@@ -126,21 +128,22 @@ def __init__(self, address):
 
         self.address = address
 
-    def forge_packet(self, dtcommands: dtprotocol.DTCommand, execute=True) -> dtprotocol.DTInstructionPacket:
+    def forge_packet(self, dtcommands: Union[List[dtprotocol.DTCommand], dtprotocol.DTCommand],
+                     execute: bool = True) -> dtprotocol.DTInstructionPacket:
         """
         Creates a packet which will be sent to the device.
 
         Args:
-            dtcommands (list): List of dtcommands.
+            dtcommands: DTCommand or list of DTCommands.
 
-            execute (bool): Sets the execute value, True by default.
+            execute: Sets the execute value, True by default.
 
         Returns:
             DTInstructionPacket: The packet created.
 
         """
         self.logger.debug("Forging packet with {} and execute set to {}".format(dtcommands, execute))
-        if type(dtcommands) == dtprotocol.DTCommand:
+        if isinstance(dtcommands, dtprotocol.DTCommand):
             dtcommands = [dtcommands]
         if execute:
             dtcommands.append(dtprotocol.DTCommand(CMD_EXECUTE))