From 97b76931c6a16dd4e99ccb254b14eb0e0840efcc Mon Sep 17 00:00:00 2001 From: Aaron LI Date: Mon, 1 Jan 2018 15:16:37 +0800 Subject: clusters/mergertree: Use utils.io for pickle dump and load --- fg21sim/extragalactic/clusters/mergertree.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) (limited to 'fg21sim/extragalactic/clusters') diff --git a/fg21sim/extragalactic/clusters/mergertree.py b/fg21sim/extragalactic/clusters/mergertree.py index 0e2a164..312fb92 100644 --- a/fg21sim/extragalactic/clusters/mergertree.py +++ b/fg21sim/extragalactic/clusters/mergertree.py @@ -6,10 +6,10 @@ Merger tree that represents the merging history of a cluster using the binary tree data structure. """ -import os -import pickle import logging +from ...utils.io import pickle_dump, pickle_load + logger = logging.getLogger(__name__) @@ -49,18 +49,12 @@ def save_mtree(mtree, outfile, clobber=False): """ Pickle the merger tree data and save to file. """ - if os.path.exists(outfile): - if clobber: - os.remove(outfile) - logger.warning("Removed existing file: {0}".format(outfile)) - else: - raise OSError("Output file already exists: {0}".format(outfile)) - pickle.dump(mtree, open(outfile, "wb")) + pickle_dump(mtree, outfile=outfile, clobber=clobber) logger.info("Saved merger tree to file: {0}".format(outfile)) def read_mtree(infile): - mtree = pickle.load(open(infile, "wb")) + mtree = pickle_load(infile) logger.info("Loaded merger tree from file: {0}".format(infile)) return mtree -- cgit v1.2.2