diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index 596c476..746b312 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -202,17 +202,16 @@ def __next__(self): except IndexError: raise StopIteration - def to_df(self, columns: list = None): # For backwards compatibility (jgray 2024) + def to_df(self, columns: list = None, force_pandas: bool = False): # For backwards compatibility (jgray 2024) """See `convert_to_dataframe`""" - return self.convert_to_dataframe(columns) + return self.convert_to_dataframe(columns, force_pandas) - def convert_to_dataframe(self, columns: list = None): + def convert_to_dataframe(self, columns: list = None, force_pandas: bool = False): """Converts the object to a pandas dataframe. - This function imports pandas and tfs-pandas, if tfs-pandas is not installed, it will only return a pandas dataframe. - Args: columns (list, optional): List of columns to include in the dataframe. Defaults to None. + force_pandas (bool, optional): If True, always use pandas.DataFrame. Defaults to False. Returns: pandas.DataFrame or tfs.TfsDataFrame: The dataframe containing the object's data. @@ -230,7 +229,10 @@ def convert_to_dataframe(self, columns: list = None): # If tfs is available, use the headers attribute DataFrame, hdr_attr = tfs.TfsDataFrame, "headers" except ImportError: - # Otherwise, use the pandas dataframe and attrs attribute + force_pandas = True + + if force_pandas: + # If pandas is the only option, use pandas DataFrame, with the attrs attribute DataFrame, hdr_attr = pd.DataFrame, "attrs" py_name, obj_name = self._mad.py_name, self._name diff --git a/tests/test_object_wrapping.py b/tests/test_object_wrapping.py index 3908940..8d8d1cf 100644 --- a/tests/test_object_wrapping.py +++ b/tests/test_object_wrapping.py @@ -4,7 +4,7 @@ import unittest import numpy as np -import pandas +import pandas as pd import tfs from pymadng import MAD @@ -369,7 +369,7 @@ def test_history(self): class TestDataFrame(unittest.TestCase): - def generalDataFrame(self, headers, DataFrame): + def generalDataFrame(self, headers, DataFrame, force_pandas=False): mad = MAD() mad.send(""" test = mtable{ @@ -394,7 +394,7 @@ def generalDataFrame(self, headers, DataFrame): test:addcol("generator", \\ri, m -> m:getcol("number")[ri] + 1i * m:getcol("number")[ri]) test:write("test") """) - df = mad.test.to_df() + df = mad.test.to_df(force_pandas=force_pandas) self.assertTrue(isinstance(df, DataFrame)) header = getattr(df, headers) self.assertEqual(header["name"], "test") @@ -431,9 +431,12 @@ def test_tfsDataFrame(self): def test_pandasDataFrame(self): sys.modules["tfs"] = None # Remove tfs-pandas - self.generalDataFrame("attrs", pandas.DataFrame) + self.generalDataFrame("attrs", pd.DataFrame) del sys.modules["tfs"] + def test_tfsDataFrame_force_pandas(self): + self.generalDataFrame("attrs", pd.DataFrame, force_pandas=True) + def test_failure(self): with MAD() as mad: mad.send("""