Skip to content

Add option to force pandas #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 5, 2025
Merged

Add option to force pandas #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 8 additions & 6 deletions src/pymadng/madp_classes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down
11 changes: 7 additions & 4 deletions tests/test_object_wrapping.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import unittest

import numpy as np
import pandas
import pandas as pd
import tfs

from pymadng import MAD
Expand Down Expand Up @@ -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{
Expand All @@ -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")
Expand Down Expand Up @@ -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("""
Expand Down