Add option to force pandas (#23)

Author: jgray-19

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

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("""