diff options
Diffstat (limited to 'fg21sim/utils')
| -rw-r--r-- | fg21sim/utils/io.py | 21 |
1 files 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)) |