Refactoring, type annotations and fixes

opendis/DataInputStream.py

@@ -4,45 +4,60 @@
 This uses big endian (network) format.
 """
 
+from io import BufferedIOBase
 import struct
 
+from .types import (
+ int8,
+ int16,
+ int32,
+ int64,
+ uint8,
+ uint16,
+ uint32,
+ float32,
+ float64,
+ char16,
+)
+
+
 class DataInputStream:
- def __init__(self, stream):
+ def __init__(self, stream: BufferedIOBase):
  self.stream = stream
 
- def read_boolean(self):
+ def read_boolean(self) -> bool:
  return struct.unpack('?', self.stream.read(1))[0]
 
- def read_byte(self):
+ def read_byte(self) -> int8:
  return struct.unpack('b', self.stream.read(1))[0]
 
- def read_unsigned_byte(self):
+ def read_unsigned_byte(self) -> uint8:
  return struct.unpack('B', self.stream.read(1))[0]
 
- def read_char(self):
+ def read_char(self) -> char16:
  return chr(struct.unpack('>H', self.stream.read(2))[0])
 
- def read_double(self):
+ def read_double(self) -> float64:
  return struct.unpack('>d', self.stream.read(8))[0]
 
- def read_float(self):
+ def read_float(self) -> float32:
  return struct.unpack('>f', self.stream.read(4))[0]
 
- def read_short(self):
+ def read_short(self) -> int16:
  return struct.unpack('>h', self.stream.read(2))[0]
 
- def read_unsigned_short(self):
+ def read_unsigned_short(self) -> uint16:
  return struct.unpack('>H', self.stream.read(2))[0]
 
- def read_long(self):
+ def read_long(self) -> int64:
  return struct.unpack('>q', self.stream.read(8))[0]
 
- def read_utf(self):
+ def read_utf(self) -> bytes:
  utf_length = struct.unpack('>H', self.stream.read(2))[0]
  return self.stream.read(utf_length)
 
- def read_int(self):
+ def read_int(self) -> int32:
  return struct.unpack('>i', self.stream.read(4))[0]
- 
- def read_unsigned_int(self):
+
+ def read_unsigned_int(self) -> uint32:
  return struct.unpack('>I', self.stream.read(4))[0]

opendis/DataOutputStream.py

@@ -4,46 +4,48 @@
 This uses big endian (network) format
 """
 
+from io import BufferedIOBase
 import struct
 
+
 class DataOutputStream:
- def __init__(self, stream):
+ def __init__(self, stream: BufferedIOBase):
  self.stream = stream
 
- def write_boolean(self, bool):
- self.stream.write(struct.pack('?', bool))
+ def write_boolean(self, boolean: bool) -> None:
+ self.stream.write(struct.pack('?', boolean))
 
- def write_byte(self, val):
+ def write_byte(self, val: bytes) -> None:
  self.stream.write(struct.pack('b', val))
 
- def write_unsigned_byte(self, val):
+ def write_unsigned_byte(self, val: bytes) -> None:
  self.stream.write(struct.pack('B', val))
 
- def write_char(self, val):
+ def write_char(self, val: str) -> None:
  self.stream.write(struct.pack('>H', ord(val)))
 
- def write_double(self, val):
+ def write_double(self, val: float) -> None:
  self.stream.write(struct.pack('>d', val))
 
- def write_float(self, val):
+ def write_float(self, val: float) -> None:
  self.stream.write(struct.pack('>f', val))
 
- def write_short(self, val):
+ def write_short(self, val: int) -> None:
  self.stream.write(struct.pack('>h', val))
 
- def write_unsigned_short(self, val):
+ def write_unsigned_short(self, val: int) -> None:
  self.stream.write(struct.pack('>H', val))
 
- def write_long(self, val):
+ def write_long(self, val: int) -> None:
  self.stream.write(struct.pack('>q', val))
 
- def write_utf(self, string):
+ def write_utf(self, string: bytes) -> None:
  self.stream.write(struct.pack('>H', len(string)))
  self.stream.write(string)
 
- def write_int(self, val):
+ def write_int(self, val: int) -> None:
  self.stream.write(struct.pack('>i', val))
 
- def write_unsigned_int(self, val):
+ def write_unsigned_int(self, val: int) -> None:
  self.stream.write(struct.pack('>I', val))
 

opendis/PduFactory.py

@@ -4,6 +4,7 @@
 from .DataInputStream import DataInputStream
 from .dis7 import *
 from io import BytesIO
+from os import PathLike
 import binascii
 import io
 
@@ -62,7 +63,7 @@
  }
 
 
-def getPdu(inputStream):
+def getPdu(inputStream: DataInputStream) -> PduSuperclass | None:
  inputStream.stream.seek(2, 0) # Skip ahead to PDU type enum field
  pduType = inputStream.read_byte()
  inputStream.stream.seek(0, 0) # Rewind to start
@@ -79,7 +80,7 @@ def getPdu(inputStream):
  return None
 
 
-def createPdu(data):
+def createPdu(data: bytes) -> PduSuperclass | None:
  """ Create a PDU of the correct type when passed an array of binary data
  input: a bytebuffer of DIS data
  output: a python DIS pdu instance of the correct class"""
@@ -90,10 +91,9 @@ def createPdu(data):
  return getPdu(inputStream)
 
 
-def createPduFromFilePath(filePath):
+def createPduFromFilePath(filePath: PathLike) -> PduSuperclass | None:
  """ Utility written for unit tests, but could have other uses too."""
- f = io.open(filePath, "rb")
- inputStream = DataInputStream(f)
- pdu = getPdu(inputStream)
- f.close()
+ with io.open(filePath, "rb") as f:
+ inputStream = DataInputStream(f)
+ pdu = getPdu(inputStream)
  return pdu

opendis/RangeCoordinates.py

@@ -46,12 +46,12 @@
 from numpy import array, dot, identity
 
 
-def deg2rad(deg):
+def deg2rad(deg: float):
  """Converts degrees to radians"""
  return deg * pi / 180
 
 
-def rad2deg(rad):
+def rad2deg(rad: float):
  """Converts radians to degrees"""
  return rad * 180 / pi