Merge pull request #74 from ngjunsiang/qol-stream-improvements

QoL improvements for data stream handling

Author: leif81

opendis/DataInputStream.py

@@ -15,8 +15,10 @@
     uint8,
     uint16,
     uint32,
+    uint64,
     float32,
     float64,
+    char8,
     char16,
 )
 
@@ -61,3 +63,41 @@ def read_int(self) -> int32:
 
     def read_unsigned_int(self) -> uint32:
         return struct.unpack('>I', self.stream.read(4))[0]
+
+    def read_bytes(self, n: int) -> bytes:
+        """Read n bytes from the stream."""
+        return self.stream.read(n)
+
+    # Aliases for convenience
+    def read_char8(self) -> char8:
+        return char8(self.read_char())
+
+    def read_float32(self) -> float32:
+        return float32(self.read_float())
+
+    def read_float64(self) -> float64:
+        return float64(self.read_double())
+
+    def read_int8(self) -> int8:
+        return int8(self.read_byte())
+
+    def read_int16(self) -> int16:
+        return int16(self.read_short())
+
+    def read_int32(self) -> int32:
+        return int32(self.read_int())
+
+    def read_int64(self) -> int64:
+        return int64(self.read_long())
+
+    def read_uint8(self) -> uint8:
+        return uint8(self.read_unsigned_byte())
+
+    def read_uint16(self) -> uint16:
+        return uint16(self.read_unsigned_short())
+
+    def read_uint32(self) -> uint32:
+        return uint32(self.read_unsigned_int())
+
+    def read_uint64(self) -> uint64:
+        return uint64(self.read_long())

opendis/DataOutputStream.py

@@ -7,6 +7,22 @@
 from io import BufferedIOBase
 import struct
 
+from .types import (
+    int8,
+    int16,
+    int32,
+    int64,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+    float32,
+    float64,
+    char8,
+    char16,
+)
+
+
 class DataOutputStream:
     def __init__(self, stream: BufferedIOBase):
         self.stream = stream
@@ -48,3 +64,40 @@ def write_int(self, val: int) -> None:
     def write_unsigned_int(self, val: int) -> None:
         self.stream.write(struct.pack('>I', val))
 
+    def write_bytes(self, val: bytes) -> None:
+        """Write bytes to the stream."""
+        self.stream.write(val)
+
+    # Aliases for convenience
+    def write_char8(self, val: char8) -> None:
+        self.write_byte(ord(val))
+
+    def write_float32(self, val: float32) -> None:
+        self.write_float(val)
+
+    def write_float64(self, val: float64) -> None:
+        self.write_double(val)
+
+    def write_int8(self, val: int8) -> None:
+        self.write_byte(val)
+
+    def write_int16(self, val: int16) -> None:
+        self.write_short(val)
+
+    def write_int32(self, val: int32) -> None:
+        self.write_int(val)
+
+    def write_int64(self, val: int64) -> None:
+        self.write_long(val)
+
+    def write_uint8(self, val: uint8) -> None:
+        self.write_unsigned_byte(val)
+
+    def write_uint16(self, val: uint16) -> None:
+        self.write_unsigned_short(val)
+
+    def write_uint32(self, val: uint32) -> None:
+        self.write_unsigned_int(val)
+
+    def write_uint64(self, val: uint64) -> None:
+        self.stream.write(struct.pack('>Q', val))

opendis/PduFactory.py

@@ -1,12 +1,12 @@
 __author__ = "mcgredo"
 __date__ = "$Jun 25, 2015 11:31:42 AM$"
 
-from .DataInputStream import DataInputStream
-from .dis7 import *
-from io import BytesIO
-from os import PathLike
-import binascii
 import io
+from os import PathLike
+
+from .dis7 import *
+from .stream import DataInputStream
+
 
 PduTypeDecoders = {
        1 : EntityStatePdu
@@ -86,7 +86,7 @@ def createPdu(data: bytes) -> PduSuperclass | None:
         input: a bytebuffer of DIS data
         output: a python DIS pdu instance of the correct class"""
 
-    memoryStream = BytesIO(data)
+    memoryStream = io.BytesIO(data)
     inputStream = DataInputStream(memoryStream)
 
     return getPdu(inputStream)

opendis/stream.py

@@ -0,0 +1,9 @@
+"""Stream utilities
+
+This module provides utilities for handling streams in the OpenDIS library.
+"""
+
+__all__ = ['DataInputStream', 'DataOutputStream']
+
+from .DataInputStream import DataInputStream
+from .DataOutputStream import DataOutputStream