Add type annotations

Author: basilfx

tests/test_tinylink.py

@@ -7,25 +7,29 @@ class DummyHandle(object):
     Dummy handler, so the TinyLink class can exchange data with itself.
     """
 
-    def __init__(self):
+    stream: bytes
+    index: int
+    length: int
+
+    def __init__(self) -> None:
         self.stream = bytearray()
         self.index = 0
         self.length = 0
 
-    def write(self, data):
+    def read(self, size: int) -> bytes:
+        data = self.stream[self.index:min(self.length, self.index + size)]
+        self.index += len(data)
+
+        # Return data.
+        return bytes(data)
+
+    def write(self, data: bytes) -> int:
         self.stream.extend(data)
         self.length += len(data)
 
-        # Return number of bytes written
+        # Return number of bytes written.
         return len(data)
 
-    def read(self, count):
-        data = self.stream[self.index:min(self.length, self.index + count)]
-        self.index += len(data)
-
-        # Return data
-        return bytes(data)
-
 
 class TinyLinkTest(unittest.TestCase):
     """

tinylink/cli.py

@@ -3,6 +3,7 @@
 import time
 import select
 import struct
+from typing import Optional
 import tinylink
 import argparse
 
@@ -14,20 +15,12 @@
     serial = None
 
 
-def run():
-    """
-    Entry point for console script.
-    """
-
-    sys.exit(main())
-
-
-def parse_arguments():
+def parse_arguments(argv: list[str]) -> argparse.Namespace:
     """
     Create and parse command line arguments.
     """
 
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(prog=argv[0])
 
     # Add options.
     parser.add_argument("port", type=str, help="serial port")
@@ -40,10 +33,10 @@ def parse_arguments():
         help="maximum length of frame")
 
     # Parse command line.
-    return parser.parse_args(), parser
+    return parser.parse_args(argv[1:])
 
 
-def dump(prefix, data):
+def dump(prefix: str, data: bytes) -> str:
     """
     Dump data as two hex columns.
     """
@@ -72,7 +65,7 @@ def dump(prefix, data):
     return "\n".join(result)
 
 
-def process_link(link):
+def process_link(link: tinylink.TinyLink) -> None:
     """
     Process incoming link data.
     """
@@ -88,7 +81,7 @@ def process_link(link):
         sys.stdout.write(dump("<<<", frame.data) + "\n\n")
 
 
-def process_stdin(link):
+def process_stdin(link: tinylink.TinyLink) -> Optional[bool]:
     """
     Process stdin commands.
     """
@@ -137,9 +130,9 @@ def process_stdin(link):
                             link.endianness + pack, int(item, 0))
                     except ValueError:
                         # Assume it is a byte string.
-                        item = item.encode("ascii")
+                        item_bytes = item.encode("ascii")
                         value = struct.pack(
-                            link.endianness + str(len(item)) + "s", item)
+                            link.endianness + str(len(item_bytes)) + "s", item_bytes)
 
                 # Concat to frame.
                 frame.data = (frame.data or bytes()) + value
@@ -162,7 +155,15 @@ def process_stdin(link):
             return
 
 
-def main():
+def run() -> None:
+    """
+    Entry point for console script.
+    """
+
+    sys.exit(main(sys.argv))
+
+
+def main(argv: list[str]) -> int:
     """
     Main entry point.
     """
@@ -174,7 +175,7 @@ def main():
         return 1
 
     # Parse arguments.
-    arguments, parser = parse_arguments()
+    arguments = parse_arguments(argv)
 
     if arguments.endianness == "little":
         endianness = tinylink.LITTLE_ENDIAN

tinylink/link.py

@@ -1,15 +1,38 @@
+from typing import Protocol
+
 from . import consts
 from . import utils
 
 import struct
 
 
-class Frame(object):
+class Handle(Protocol):
+    """
+    Protocol for a handler.
+    """
+
+    def read(self, size: int) -> bytes:
+        """"
+        Read up to  `size` bytes.
+        """
+
+    def write(self, data: bytes) -> int:
+        """
+        Write data and return the number of bytes written.
+        """
+
+
+class Frame:
     """
     Represents a frame.
     """
 
-    def __init__(self, data=None, flags=consts.FLAG_NONE, damaged=False):
+    data: bytes
+    flags: int
+    damaged: int
+
+    def __init__(self, data: bytes = None, flags: int = consts.FLAG_NONE,
+                 damaged: bool = False) -> None:
         if data is not None:
             if type(data) is not bytes:
                 raise ValueError("Provided data must be encoded as bytes.")
@@ -20,12 +43,12 @@ def __init__(self, data=None, flags=consts.FLAG_NONE, damaged=False):
         self.flags = flags
         self.damaged = damaged
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "%s(%s, flags=%d, damaged=%s)" % (
             self.__class__.__name__, repr(self.data), self.flags, self.damaged)
 
 
-class TinyLink(object):
+class TinyLink:
     """
     TinyLink state machine for streaming communication with low-speed embedded
     applications that only use RX/TX. Every message is encapsulated in a frame.
@@ -39,8 +62,14 @@ class TinyLink(object):
     It does not provide error correction and the bytes are not aligned.
     """
 
-    def __init__(self, handle, endianness=consts.LITTLE_ENDIAN,
-                 max_length=2**(consts.LEN_LENGTH * 8), ignore_damaged=False):
+    handle: Handle
+    endianness: str
+    max_length: int
+    ignore_damaged: bool
+
+    def __init__(self, handle: Handle, endianness: str = consts.LITTLE_ENDIAN,
+                 max_length: int = 2**(consts.LEN_LENGTH * 8),
+                 ignore_damaged: bool = False) -> None:
         """
         Construct a new TinyLink state machine. A state machine takes a handle,
         which provides a `read` and `write` method.
@@ -75,7 +104,7 @@ def __init__(self, handle, endianness=consts.LITTLE_ENDIAN,
         # Python 2 does not allow unpack from bytearray, but Python 3.
         self.buffer = self.stream
 
-    def write_frame(self, frame):
+    def write_frame(self, frame: Frame) -> int:
         """
         Write a frame via the handle.
         """
@@ -105,14 +134,14 @@ def write_frame(self, frame):
         # Write to file.
         return self.handle.write(result)
 
-    def write(self, data, flags=consts.FLAG_NONE):
+    def write(self, data: bytes, flags: int = consts.FLAG_NONE) -> int:
         """
         Shorthand for `write_frame(Frame(data, flags=flags))`.
         """
 
         return self.write_frame(Frame(data, flags=flags))
 
-    def read(self, limit=1):
+    def read(self, limit: int = 1) -> list[Frame]:
         """
         Read up to `limit` bytes from the handle and process it. Returns a list
         of received frames, if any.