ucftar downsampler from lsdf

python/programs/ucftar_downsampler

@@ -7,22 +7,32 @@ import numpy as np
 import ucf
 
 parser = argparse.ArgumentParser(description='Reads an ucf.tar archive, downsamples it and saves it to a new ucf.tar archive. Can be used as a pipe.')
+parser.add_argument("-c", "--chunks", metavar='file',nargs='?', type=int, default=None, help="number of chunks [default: not chunked]", action="store")
 parser.add_argument("-i", "--infile", metavar='file',nargs='?', default=None, help="name of the input file [default: stdin]", action="store")
 parser.add_argument("-o", "--outfile", metavar='file',nargs='?', default=None, help="name of the output file [default: stdout]", action="store")
 parser.add_argument("-n", "--nskip", metavar='N',nargs='?', type=int, default=2, help="keep every Nth grid point [default: 2]", action="store")
 parser.add_argument("-sp", "--single-precision", help="output data in single-precision? [default: False]", action="store_true")
 args = parser.parse_args()
 
+nchunk   = args.chunks
 nskip    = args.nskip
 file_in  = args.infile
 file_out = args.outfile
 saveSinglePrecision = args.single_precision
 
-if file_in is None:
-  istream = tarfile.open(fileobj=sys.stdin.buffer,mode='r|',bufsize=512*1024**2,ignore_zeros=True)
+if nchunk is None:
+  if file_in is None:
+    istream = tarfile.open(fileobj=sys.stdin.buffer,mode='r|',bufsize=512*1024**2,ignore_zeros=True)
+  else:
+    filehandle_in = open(file_in,'rb')
+    istream = tarfile.open(fileobj=filehandle_in,mode='r')
 else:
-  filehandle_in = open(file_in,'rb')
-  istream = tarfile.open(fileobj=filehandle_in,mode='r')
+  if file_in is None:
+    raise ValueError('Chunked mode only works with input file, not stream')
+  else:
+    ichunk = 0
+    filehandle_in = open(file_in+'.{:d}'.format(ichunk),'rb')
+    istream = tarfile.open(fileobj=filehandle_in,mode='r')
 
 if file_out is None:
   ostream = tarfile.open(fileobj=sys.stdout.buffer,mode='w|',bufsize=512*1024**2,pax_headers=tarfile.USTAR_FORMAT)
@@ -33,6 +43,17 @@ else:
 while True:
     iinfo = istream.next()
     if iinfo is None:
+      if nchunk is not None:
+        if ichunk+1<nchunk:
+          istream.close()
+          filehandle_in.close()
+          ichunk = ichunk+1
+          filehandle_in = open(file_in+'.{:d}'.format(ichunk),'rb')
+          istream = tarfile.open(fileobj=filehandle_in,mode='r')
+          iinfo = istream.next()
+        else:
+          break
+      else:
         break
     print(iinfo.name,file=sys.stderr)