aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorAaron LI <aly@aaronly.me>2017-08-01 00:58:55 +0800
committerAaron LI <aly@aaronly.me>2017-08-01 00:58:55 +0800
commitadc667c206b8c3438be8a94f3a839edc3a077ce0 (patch)
tree144bfb77b55eac4eec9946acaea88d8ed43f9009
parent7c09bf216640506c04bb2ad1cafb58f3406fe64c (diff)
downloadfg21sim-adc667c206b8c3438be8a94f3a839edc3a077ce0.tar.bz2
utils/io.py: Add functions "pickle_dump()" and "pickle_load()"
-rw-r--r--fg21sim/utils/io.py31
1 files changed, 31 insertions, 0 deletions
diff --git a/fg21sim/utils/io.py b/fg21sim/utils/io.py
index ccecf21..c19ef47 100644
--- a/fg21sim/utils/io.py
+++ b/fg21sim/utils/io.py
@@ -7,6 +7,7 @@ Input/output utilities.
import os
import logging
+import pickle
from datetime import datetime
import pandas as pd
@@ -77,3 +78,33 @@ def dataframe_to_csv(df, outfile, comment=None, clobber=False):
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))
+
+
+def pickle_dump(obj, outfile, clobber=False):
+ """
+ Dump the given object into the output file using ``pickle.dump()``.
+
+ NOTE
+ ----
+ The dumped output file is in binary format, and can be loaded back
+ using ``pickle.load()``.
+
+ Example
+ -------
+ >>> a = [1, 2, 3]
+ >>> pickle.dump(a, file=open("a.pkl", "wb"))
+ >>> b = pickle.load(open("a.pkl", "rb))
+ >>> a == b
+ True
+
+ Parameters
+ ----------
+ """
+ _create_dir(outfile)
+ _check_existence(outfile, clobber=clobber, remove=True)
+ pickle.dump(obj, file=open(outfile, "wb"))
+ logger.info("Pickled data to file: %s" % outfile)
+
+
+def pickle_load(infile):
+ return pickle.load(open(infile, "rb"))