Keep file handles an write to them while iterating through data

dasscripts/dasconv.py

@@ -21,6 +21,7 @@ import sys
 import logging
 import pprint
 import datetime
+import numpy as np
 from obspy import Trace
 from obspy import Stream
 from .tdms import TDMS
@@ -108,9 +109,14 @@ def main():
               chstart=args.chstart, chstop=args.chstop, chstep=args.chstep,
               starttime=dtstart, endtime=dtend, decimate=args.decimate)
 
+
+    filehandles = dict()
+
     with td:
         curstream = None
         for data in td:
+            print(data[1], data[0])
+
             if args.metadata:
                 if isinstance(data, dict):
                     pprint.pprint(data)
@@ -120,10 +126,14 @@ def main():
                 if curstream != '%s.%s.%s.%s' % (data[1]['network'], data[1]['station'], data[1]['location'], data[1]['channel']):
                     # Save the previous Stream completely
                     if curstream is not None:
-                        st = Stream([tr0])
+                        # st = Stream([tr0])
                         logs.info('Storing channel %s' % curstream)
-                        # FIXME This overwrites the previous chunk of data!
-                        st.write('%s.mseed' % curstream, format='MSEED')
+                        try:
+                            tr0.write(filehandles[curstream], format='MSEED')
+                        except KeyError:
+                            filehandles[curstream] = open('%s.mseed' % curstream, 'wb')
+                            print(1, filehandles)
+                            tr0.write(filehandles[curstream], format='MSEED')
 
                     # Update which stream is being processed
                     curstream = '%s.%s.%s.%s' % (data[1]['network'], data[1]['station'], data[1]['location'], data[1]['channel'])
@@ -131,12 +141,22 @@ def main():
                     tr0 = Trace(data=data[0], header=data[1])
                 else:
                     tr0 += Trace(data=data[0], header=data[1])
+                    # print(data[1], data[0])
         else:
             if not args.metadata:
-                st = Stream([tr0])
-                logs.info('Storing channel %s' % curstream)
-                # FIXME This overwrites the previous chunk of data!
-                st.write('%s.mseed' % curstream, format='MSEED')
+                # st = Stream([tr0])
+                logs.info('Storing last part of channel %s' % curstream)
+                # st.write('%s-%08d.mseed' % (curstream, fnumber), format='MSEED')
+                print(2, filehandles)
+                try:
+                    tr0.write(filehandles[curstream], format='MSEED')
+                except KeyError:
+                    filehandles[curstream] = open('%s.mseed' % curstream, 'wb')
+                    print(3, filehandles)
+                    # tr0.write(filehandles[curstream], format='MSEED')
+                    if isinstance(tr0.data, np.ma.masked_array):
+                        tr0.data = tr0.data.filled()
+                    tr0.write(filehandles[curstream], format='MSEED')
 
 
 if __name__ == '__main__':