From 52d8ab190765ca543ea8978368392658e94b9789 Mon Sep 17 00:00:00 2001 From: Daniel Weindl Date: Wed, 26 Jul 2023 15:11:11 +0200 Subject: [PATCH] Handle DataFrame/None as input to `get_visualization_df` (#221) For coherence with other `get_XXX_df` functions. --- petab/core.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/petab/core.py b/petab/core.py index 30c4c017..ffd5ea86 100644 --- a/petab/core.py +++ b/petab/core.py @@ -54,15 +54,25 @@ def write_simulation_df(df: pd.DataFrame, filename: Union[str, Path]) -> None: df.to_csv(filename, sep='\t', index=False) -def get_visualization_df(visualization_file: Union[str, Path]) -> pd.DataFrame: +def get_visualization_df( + visualization_file: Union[str, Path, pd.DataFrame, None] +) -> Union[pd.DataFrame, None]: """Read PEtab visualization table Arguments: - visualization_file: URL or filename of PEtab visualization table + visualization_file: + URL or filename of PEtab visualization table to read from, + or a DataFrame or None that will be returned as is. Returns: Visualization DataFrame """ + if visualization_file is None: + return None + + if isinstance(visualization_file, pd.DataFrame): + return visualization_file + try: types = {PLOT_NAME: str} vis_spec = pd.read_csv(visualization_file, sep="\t", index_col=None,