| from typing import cast |
|
|
| import pandas as pd |
| from pandas import DataFrame as pandas_DataFrame |
|
|
| from langflow.schema.data import Data |
|
|
|
|
| class DataFrame(pandas_DataFrame): |
| """A pandas DataFrame subclass specialized for handling collections of Data objects. |
| |
| This class extends pandas.DataFrame to provide seamless integration between |
| Langflow's Data objects and pandas' powerful data manipulation capabilities. |
| |
| Args: |
| data: Input data in various formats: |
| - List[Data]: List of Data objects |
| - List[Dict]: List of dictionaries |
| - Dict: Dictionary of arrays/lists |
| - pandas.DataFrame: Existing DataFrame |
| - Any format supported by pandas.DataFrame |
| **kwargs: Additional arguments passed to pandas.DataFrame constructor |
| |
| Examples: |
| >>> # From Data objects |
| >>> dataset = DataFrame([Data(data={"name": "John"}), Data(data={"name": "Jane"})]) |
| |
| >>> # From dictionaries |
| >>> dataset = DataFrame([{"name": "John"}, {"name": "Jane"}]) |
| |
| >>> # From dictionary of lists |
| >>> dataset = DataFrame({"name": ["John", "Jane"], "age": [30, 25]}) |
| """ |
|
|
| def __init__(self, data: None | list[dict | Data] | dict | pd.DataFrame = None, **kwargs): |
| if data is None: |
| super().__init__(**kwargs) |
| return |
|
|
| if isinstance(data, list): |
| if all(isinstance(x, Data) for x in data): |
| data = [d.data for d in data if hasattr(d, "data")] |
| elif not all(isinstance(x, dict) for x in data): |
| msg = "List items must be either all Data objects or all dictionaries" |
| raise ValueError(msg) |
| kwargs["data"] = data |
| elif isinstance(data, dict | pd.DataFrame): |
| kwargs["data"] = data |
|
|
| super().__init__(**kwargs) |
|
|
| def to_data_list(self) -> list[Data]: |
| """Converts the DataFrame back to a list of Data objects.""" |
| list_of_dicts = self.to_dict(orient="records") |
| return [Data(data=row) for row in list_of_dicts] |
|
|
| def add_row(self, data: dict | Data) -> "DataFrame": |
| """Adds a single row to the dataset. |
| |
| Args: |
| data: Either a Data object or a dictionary to add as a new row |
| |
| Returns: |
| DataFrame: A new DataFrame with the added row |
| |
| Example: |
| >>> dataset = DataFrame([{"name": "John"}]) |
| >>> dataset = dataset.add_row({"name": "Jane"}) |
| """ |
| if isinstance(data, Data): |
| data = data.data |
| new_df = self._constructor([data]) |
| return cast("DataFrame", pd.concat([self, new_df], ignore_index=True)) |
|
|
| def add_rows(self, data: list[dict | Data]) -> "DataFrame": |
| """Adds multiple rows to the dataset. |
| |
| Args: |
| data: List of Data objects or dictionaries to add as new rows |
| |
| Returns: |
| DataFrame: A new DataFrame with the added rows |
| """ |
| processed_data = [] |
| for item in data: |
| if isinstance(item, Data): |
| processed_data.append(item.data) |
| else: |
| processed_data.append(item) |
| new_df = self._constructor(processed_data) |
| return cast("DataFrame", pd.concat([self, new_df], ignore_index=True)) |
|
|
| @property |
| def _constructor(self): |
| def _c(*args, **kwargs): |
| return DataFrame(*args, **kwargs).__finalize__(self) |
|
|
| return _c |
|
|
| def __bool__(self): |
| """Truth value testing for the DataFrame. |
| |
| Returns True if the DataFrame has at least one row, False otherwise. |
| """ |
| return not self.empty |
|
|
