diff --git a/bf_pulp_utils/bf_pulp_utils/psrfits_requantisation.py b/bf_pulp_utils/bf_pulp_utils/psrfits_requantisation.py index dd465785cb5ff4176e60d39ac726f273800689d8..8fdc487ffe42cdcdabc71ab0b07e68d1aa2e70cb 100644 --- a/bf_pulp_utils/bf_pulp_utils/psrfits_requantisation.py +++ b/bf_pulp_utils/bf_pulp_utils/psrfits_requantisation.py @@ -97,16 +97,21 @@ def main(): # getting current directory workdir=os.getcwd() if len(matches) != 0: + rootdir="root" + os.mkdir(rootdir) # create "root" directory in the current dir where all files will be extracted from the tar # extract .fits files from the output tarball outputtar = tarfile.open(output_tarball, "r") print ("Extracting *.fits files:") for ii in matches: print (ii) - outputtar.extract(ii, "./") + #outputtar.extract(ii, "./") + outputtar.extractall(rootdir) outputtar.close() + # remove the input tarball (as the output tarfile will have the same name) + os.system("rm -f %s" % output_tarball) + os.chdir(rootdir) # converting extracted 8-bit PSRFITS files to 2- or 4-bit PSRFITS files - outputtar = tarfile.open(output_tarball, "a") - print ("Converting to %d-bit and adding to the output tarball..." % (nbit)) + print ("Converting to %d-bit..." % (nbit)) for ii in matches: print (ii) options="-outbits %d -adjustlevels -numfiles 1" % (nbit) @@ -130,26 +135,24 @@ def main(): raise Exception # renaming new converted fitsfile to the original name os.rename("%s_0001.fits" % (outputbase), "%s.fits" % (outputbase)) - # adding new 2/4-bit PSRFITS and logfile (and readme?) to the output tarball - outputtar.add("%s.fits" % (outputbase)) + # adding new files to 'to_add' list and original file to 'to_delete' list to_add.append("%s.fits" % (outputbase)) - outputtar.add(logfile) to_add.append(logfile) - # removing the PSRFITS file from the working directory - os.remove(ii) to_delete.append(ii) - os.remove("%s.fits" % (outputbase)) + # removing the original PSRFITS file + os.remove(ii) except Exception as error: print ("=========Error has occurred: ========") print (str(error)) sys.exit(1) - outputtar.close() - # deleting original 8-bit PSRFITS files from the output tarball - print ("Deleting original 8-bit PSRFITS files:") - for ii in matches: - print(ii) - # Python tarfile module does not allow to delete files from the archives, using system calls for now - os.system("tar -vf %s --delete %s" % (output_tarball, ii)) + + os.chdir("../") + print ("Creating new tarball %s..." % (output_tarball)) + with tarfile.open(output_tarball, 'w') as outputtar: + for filename in os.listdir(rootdir): + fpath = os.path.join(rootdir, filename) + outputtar.add(fpath, arcname=filename) + os.system("rm -rf %s" % (rootdir)) else: print ("No PSRFITS files found. Nothing to do. Exiting...")