aboutsummaryrefslogtreecommitdiffstats
path: root/fg21sim
diff options
context:
space:
mode:
authorAaron LI <aly@aaronly.me>2017-07-19 22:28:05 +0800
committerAaron LI <aly@aaronly.me>2017-07-19 22:28:05 +0800
commit673ba9d15831505e8116fde4c3bb86de918ee462 (patch)
treed979d48a65665005e2b791a6711cc6a467d9435c /fg21sim
parent7d1b2c32ab823a06bf734b24e00daa3c6f3fd03f (diff)
downloadfg21sim-673ba9d15831505e8116fde4c3bb86de918ee462.tar.bz2
dataframe_to_csv(): Add parameter "comment"
Also write the writer and datetime to header by default. Signed-off-by: Aaron LI <aly@aaronly.me>
Diffstat (limited to 'fg21sim')
-rw-r--r--fg21sim/utils/io.py21
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))