Add an object TbTData reader

turn_by_turn/io.py

@@ -10,7 +10,7 @@
 from pathlib import Path
 from typing import Union, Any
 
-from turn_by_turn import ascii, doros, esrf, iota, lhc, ptc, sps, trackone, madng
+from turn_by_turn import ascii, doros, esrf, iota, lhc, ptc, sps, trackone, madng, tbt_data
 from turn_by_turn.ascii import write_ascii
 from turn_by_turn.errors import DataTypeError
 from turn_by_turn.structures import TbtData
@@ -30,26 +30,28 @@
     trackone=trackone,
     ascii=ascii,
     madng=madng,
+    tbt_data=tbt_data,
 )
 WRITERS = ("lhc", "sps", "doros", "doros_positions", "doros_oscillations", "ascii", "madng")  # implemented writers
 
 write_lhc_ascii = write_ascii  # Backwards compatibility <0.4
 
 
-def read_tbt(file_path: Union[str, Path], datatype: str = "lhc") -> TbtData:
+def read_tbt(file_path: Union[str, Path, TbtData], datatype: str = "lhc") -> TbtData:
     """
     Calls the appropriate loader for the provided matrices type and returns a ``TbtData`` object of the
     loaded matrices.
 
     Args:
-        file_path (Union[str, Path]): path to a file containing TbtData.
+        file_path (Union[str, Path, TbtData]): path to a file containing TbtData.
         datatype (str): type of matrices in the file, determines the reader to use. Case-insensitive,
             defaults to ``lhc``.
 
     Returns:
         A ``TbtData`` object with the loaded matrices.
     """
-    file_path = Path(file_path)
+    if not isinstance(file_path, TbtData):
+        file_path = Path(file_path)
     LOGGER.info(f"Loading turn-by-turn matrices from '{file_path}'")
     try:
         module = TBT_MODULES[datatype.lower()]

turn_by_turn/tbt_data.py

@@ -0,0 +1,45 @@
+"""
+Generic TbT Data
+---
+
+Data handling for turn-by-turn measurement data.
+Can be used with in-memory tracking data so not to have to write .sdds files.
+
+"""
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, Union
+
+import numpy as np
+import pandas as pd
+import sdds
+from dateutil import tz
+
+from turn_by_turn.ascii import is_ascii_file, read_tbt as read_ascii
+from turn_by_turn.constants import PLANES, PLANE_TO_NUM
+from turn_by_turn.structures import TbtData, TransverseData
+from turn_by_turn.utils import matrices_to_array
+
+LOGGER = logging.getLogger()
+
+def read_tbt(data: TbtData) -> TbtData:
+    """
+    Reads turn-by-turn data as a TbtData format.
+
+    Args:
+        data (Union[TbtData]): measurement data.
+
+    Returns:
+        The same ``TbTData`` object that was loaded.
+    """
+
+    return data
+
+
+def write_tbt(*args, **kwargs) -> None:
+    """
+    Not implemented, as it would not make sense to write an object. Use other functions to write to
+    the desired format.
+    """
+    raise NotImplementedError("This function is not implemented.")