make trf2pandas be able to handle directly passing the binary file contents
"""Decodes trf file.
"""
+import os # pylint: disable = unused-import
+from typing import Any, BinaryIO, Tuple, Union, cast # pylint: disable = unused-import
+
from pymedphys._imports import pandas as pd
from .header import Header, decode_header
from .partition import split_into_header_table
from .table import decode_trf_table
+path_or_binary_file = Union[BinaryIO, "os.PathLike[Any]"]
+
+
+def trf2pandas(trf: path_or_binary_file) -> Tuple["pd.DataFrame", "pd.DataFrame"]:
+ binary_file_trf = cast(BinaryIO, trf)
+ path_like_trf = cast("os.PathLike[Any]", trf)
-def trf2pandas(filepath):
- with open(filepath, "rb") as file:
- trf_contents = file.read()
+ try:
+ binary_file_trf.seek(0)
+ trf_contents = binary_file_trf.read()
+ except AttributeError:
+ with open(path_like_trf, "rb") as f:
+ trf_contents = f.read()
trf_header_contents, trf_table_contents = split_into_header_table(trf_contents)
