fixup consolidate dispatch registers

Author: medley56

space_packet_parser/generators/utils.py

@@ -2,40 +2,39 @@
 
 import datetime as dt
 import io
+import logging
 import socket
 import time
 from functools import singledispatch
 from os import PathLike
 from pathlib import Path
 from typing import Optional, Union
 
-from space_packet_parser.common import logger
+logger = logging.getLogger(__name__)
 
 
 @singledispatch
 def _read_packet_file(packet_file) -> Union[bytes, io.BufferedIOBase, io.RawIOBase]:
-    """Read a packet file or file-like object and return a bytes-reader suitable for passing to a generator.
+    """Read a packet file or file-like object and return an object suitable for passing to a generator.
+
+    Specifically this function prepares the input for use with _setup_binary_reader.
 
     Parameters
     ----------
     packet_file : Union[str, Path, PathLike, BinaryIO, bytes]
 
     Notes
     -----
-    This function handles pathlike objects, turning them into objects compatible with generators.
+    This function handles strings and pathlike objects but it reads the full file into memory for the generator.
     This alleviates the need for generators to internally handle opening and closing files.
+    For a more memory efficient approach, pass an opened file object.
     """
     raise OSError(f"Unable to open and read packet_file type: {type(packet_file)}")
 
 
-@_read_packet_file.register
-def _(packet_file: io.BufferedIOBase) -> io.BufferedIOBase:
-    """File-like object, this can be passed directly to a generator."""
-    return packet_file
-
-
-@_read_packet_file.register
-def _(packet_file: io.RawIOBase) -> io.RawIOBase:
+@_read_packet_file.register(io.BufferedIOBase)
+@_read_packet_file.register(io.RawIOBase)
+def _(packet_file: Union[io.BufferedIOBase, io.RawIOBase]) -> Union[io.BufferedIOBase, io.RawIOBase]:
     """File-like object, this can be passed directly to a generator."""
     return packet_file
 
@@ -48,12 +47,7 @@ def _(packet_file: bytes) -> bytes:
 
 @_read_packet_file.register
 def _(packet_file: str) -> bytes:
-    """String file path, open and read bytes.
-
-    Notes
-    -----
-    This will not work for PathLike objects that are not strings, use the Path or Path
-    """
+    """String file path, open and read bytes."""
     with open(packet_file, "rb") as f:
         return f.read()
 
@@ -83,7 +77,7 @@ def _(packet_file: PathLike) -> bytes:
 
 
 @singledispatch
-def _setup_binary_reader(binary_data, buffer_read_size_bytes=None):
+def _setup_binary_reader(binary_data, buffer_read_size_bytes=None) -> tuple:
     """Helper to set up reading from binary_data (file, socket, bytes) for a packet generator.
 
     Parameters
@@ -104,24 +98,9 @@ def _setup_binary_reader(binary_data, buffer_read_size_bytes=None):
     raise OSError(f"Unrecognized data source: {binary_data}")
 
 
-@_setup_binary_reader.register
-def _(binary_data: io.BufferedIOBase, buffer_read_size_bytes=None):
-    """Set up a binary reader from a file-like object."""
-    read_buffer = b""
-    if buffer_read_size_bytes is None:
-        # Default to a full read of the file
-        buffer_read_size_bytes = -1
-    total_length_bytes = binary_data.seek(0, io.SEEK_END)
-    binary_data.seek(0, 0)
-    read_bytes_from_source = binary_data.read
-    logger.info(
-        f"Creating packet generator from a filelike object, {binary_data}. Total length is {total_length_bytes} bytes"
-    )
-    return read_buffer, total_length_bytes, read_bytes_from_source, buffer_read_size_bytes
-
-
-@_setup_binary_reader.register
-def _(binary_data: io.RawIOBase, buffer_read_size_bytes=None):
+@_setup_binary_reader.register(io.BufferedIOBase)
+@_setup_binary_reader.register(io.RawIOBase)
+def _(binary_data: Union[io.BufferedIOBase, io.RawIOBase], buffer_read_size_bytes=None) -> tuple:
     """Set up a binary reader from a file-like object."""
     read_buffer = b""
     if buffer_read_size_bytes is None:
@@ -137,7 +116,7 @@ def _(binary_data: io.RawIOBase, buffer_read_size_bytes=None):
 
 
 @_setup_binary_reader.register
-def _(binary_data: socket.socket, buffer_read_size_bytes=None):
+def _(binary_data: socket.socket, buffer_read_size_bytes=None) -> tuple:
     """Set up a binary reader from a socket object."""
     read_buffer = b""
     total_length_bytes = None  # We don't know how long it is
@@ -150,7 +129,7 @@ def _(binary_data: socket.socket, buffer_read_size_bytes=None):
 
 
 @_setup_binary_reader.register
-def _(binary_data: bytes, buffer_read_size_bytes=None):
+def _(binary_data: bytes, buffer_read_size_bytes=None) -> tuple:
     """Set up a binary reader from a bytes object."""
     read_buffer = b""
     read_buffer = binary_data