dataframe_to_csv(): Add parameter "comment"

Also write the writer and datetime to header by default.

Signed-off-by: Aaron LI <aly@aaronly.me>

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))