diff options
Diffstat (limited to 'fg21sim/utils')
| -rw-r--r-- | fg21sim/utils/io.py | 37 |
1 files changed, 35 insertions, 2 deletions
diff --git a/fg21sim/utils/io.py b/fg21sim/utils/io.py index a1a5000..2449ca5 100644 --- a/fg21sim/utils/io.py +++ b/fg21sim/utils/io.py @@ -93,13 +93,14 @@ def _check_existence(filepath, clobber=False, remove=False): def dataframe_to_csv(df, outfile, comment=None, clobber=False): """ - Save the given Pandas DataFrame into a CSV text file. + Save the given Pandas DataFrame into a CSV text file with comments + prepended at the file head. Parameters ---------- df : `~pandas.DataFrame` The DataFrame to be saved to the CSV text file. - outfile : string + outfile : str The path to the output CSV file. comment : list[str], optional A list of comments to be prepended to the output CSV file header. @@ -128,6 +129,38 @@ def dataframe_to_csv(df, outfile, comment=None, clobber=False): logger.info("Wrote DataFrame to CSV file: {0}".format(outfile)) +def csv_to_dataframe(infile): + """ + Read the given CSV file as a Pandas DataFrame, with head comments + also considered and returned. + + Parameters + ---------- + infile : str + The path to the input CSV file. + + Returns + df : `~pandas.DataFrame` + The DataFrame read from the CSV text file. + comment : list[str] + A list of comments read from the lines prefixing with ``#`` + at the CSV file header. + The prefix ``#`` is striped. + """ + comments = [] + for line in open(infile): + line = line.strip() + if line == "": + continue + elif line[0] == "#": + comments.append(line.lstrip("# ")) + else: + break + + df = pd.read_csv(infile, comment="#") + return (df, comments) + + def pickle_dump(obj, outfile, clobber=False): """ Dump the given object into the output file using ``pickle.dump()``. |