From 673ba9d15831505e8116fde4c3bb86de918ee462 Mon Sep 17 00:00:00 2001 From: Aaron LI Date: Wed, 19 Jul 2017 22:28:05 +0800 Subject: dataframe_to_csv(): Add parameter "comment" Also write the writer and datetime to header by default. Signed-off-by: Aaron LI --- fg21sim/utils/io.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/fg21sim/utils/io.py b/fg21sim/utils/io.py index 4d9999e..be38032 100644 --- a/fg21sim/utils/io.py +++ b/fg21sim/utils/io.py @@ -7,6 +7,7 @@ Input/output utilities. import os import logging +from datetime import datetime import pandas as pd @@ -14,7 +15,7 @@ import pandas as pd logger = logging.getLogger(__name__) -def dataframe_to_csv(df, outfile, clobber=False): +def dataframe_to_csv(df, outfile, comment=None, clobber=False): """ Save the given Pandas DataFrame into a CSV text file. @@ -24,6 +25,9 @@ def dataframe_to_csv(df, outfile, clobber=False): The DataFrame to be saved to the CSV text file. outfile : string The path to the output CSV file. + comment : list[str], optional + A list of comments to be prepended to the output CSV file header. + The prefix ``#`` is not required and will be automatically added. clobber : bool, optional Whether overwrite the existing output file? Default: False @@ -39,9 +43,18 @@ def dataframe_to_csv(df, outfile, clobber=False): if os.path.exists(outfile): if clobber: - logger.warning("Remove existing file: {0}".format(outfile)) + logger.warning("Removed existing file: {0}".format(outfile)) os.remove(outfile) else: raise OSError("Output file exists: {0}".format(outfile)) - df.to_csv(outfile, header=True, index=False) - logger.info("Saved DataFrame to CSV file: {0}".format(outfile)) + + # Add a default header comment + if comment is None: + comment = ["by %s" % __name__, + "at %s" % datetime.now().isoformat()] + + with open(outfile, "w") as fh: + # Write header comments with ``#`` prefixed. + fh.write("".join(["# "+line.strip()+"\n" for line in comment])) + df.to_csv(fh, header=True, index=False) + logger.info("Wrote DataFrame to CSV file: {0}".format(outfile)) -- cgit v1.2.2