Skip to content

Commit

Permalink
added to_pandas method to convert CSV objects by default
Browse files Browse the repository at this point in the history
  • Loading branch information
arulmabr committed Jan 7, 2025
1 parent bc42e9d commit fa3504b
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 0 deletions.
32 changes: 32 additions & 0 deletions edsl/scenarios/FileStore.py
Original file line number Diff line number Diff line change
Expand Up @@ -327,6 +327,38 @@ def create_link(self, custom_filename=None, style=None):

return ConstructDownloadLink(self).create_link(custom_filename, style)

def to_pandas(self):
"""
Convert the file content to a pandas DataFrame if supported by the file handler.
Returns:
pandas.DataFrame: The data from the file as a DataFrame
Raises:
AttributeError: If the file type's handler doesn't support pandas conversion
"""
handler = FileMethods.get_handler(self.suffix)
if handler and hasattr(handler, "to_pandas"):
return handler(self.path).to_pandas()
raise AttributeError(
f"Converting {self.suffix} files to pandas DataFrame is not supported"
)

def __getattr__(self, name):
"""
Delegate pandas DataFrame methods to the underlying DataFrame if this is a CSV file
"""
if self.suffix == "csv":
# Get the pandas DataFrame
df = self.to_pandas()
# Check if the requested attribute exists in the DataFrame
if hasattr(df, name):
return getattr(df, name)
# If not a CSV or attribute doesn't exist in DataFrame, raise AttributeError
raise AttributeError(
f"'{self.__class__.__name__}' object has no attribute '{name}'"
)


class CSVFileStore(FileStore):
@classmethod
Expand Down
11 changes: 11 additions & 0 deletions edsl/scenarios/handlers/csv.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,3 +36,14 @@ def example(self):
with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as f:
df.to_csv(f.name, index=False)
return f.name

def to_pandas(self):
"""
Convert the CSV file to a pandas DataFrame.
Returns:
pandas.DataFrame: The data from the CSV as a DataFrame
"""
import pandas as pd

return pd.read_csv(self.path)

0 comments on commit fa3504b

Please sign in to comment.