From 160efe88a76d95e549ac77cfc47a7b7ce4dcfb84 Mon Sep 17 00:00:00 2001 From: Miguel Marques Date: Mon, 8 Mar 2021 10:25:20 +0100 Subject: [PATCH] Added the possibility of reading compressed vasprun.xml files. --- dpdata/vasp/xml.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/dpdata/vasp/xml.py b/dpdata/vasp/xml.py index d5e1fb039..6bb885947 100755 --- a/dpdata/vasp/xml.py +++ b/dpdata/vasp/xml.py @@ -75,6 +75,16 @@ def analyze (fname, type_idx_zero = False, begin = 0, step = 1) : """ can deal with broken xml file """ + + import os, bz2, gzip + fh = fname + if(isinstance(fname, str)): + fext = os.path.splitext(fname)[1] + if(fext == ".bz2"): + fh = bz2.open(fname, "r") + elif(fext == ".gz"): + fh = gzip.open(fname, "r") + all_posi = [] all_cell = [] all_ener = [] @@ -82,7 +92,7 @@ def analyze (fname, type_idx_zero = False, begin = 0, step = 1) : all_strs = [] cc = 0 try: - for event, elem in ET.iterparse(fname): + for event, elem in ET.iterparse(fh): if elem.tag == 'atominfo' : eles, types = analyze_atominfo(elem) types = np.array(types, dtype = int)