1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
5from __future__ import division, absolute_import
7import os
8import random
10from .io_common import FileLoadError, FileSaveError
11from pyrocko.util import ensuredirs
14def detect(first512):
15 lines = first512.lstrip().splitlines()
16 if len(lines) >= 2:
17 if lines[0].startswith(b'WID2 '):
18 return True
20 if lines[0].startswith(b'BEGIN GSE2'):
21 return True
23 if lines[0].startswith(b'DATA_TYPE WAVEFORM GSE2'):
24 return True
26 return False
29def iload(filename, load_data=True):
31 from . import ims
33 try:
34 with open(filename, 'rb') as f:
36 r = ims.Reader(f, load_data=load_data, version='GSE2.0',
37 dialect=None)
39 for sec in r:
40 if isinstance(sec, ims.WID2Section):
41 tr = sec.pyrocko_trace(checksum_error='warn')
42 yield tr
44 except (OSError, ims.DeserializeError) as e:
45 fle = FileLoadError(e)
46 fle.set_context('filename', filename)
47 raise fle
50def randomid():
51 return ''.join(chr(random.randint(97, 122)) for _ in range(20))
54def save(traces, filename_template, additional={}, max_open_files=10,
55 overwrite=True):
57 from pyrocko import info
58 from . import ims
60 fns = set()
61 open_files = {}
63 def close_files():
64 while open_files:
65 open_files.popitem()[1].close()
67 for tr in traces:
68 fn = tr.fill_template(filename_template, **additional)
69 if fn not in open_files:
70 if len(open_files) >= max_open_files:
71 close_files()
73 if fn not in fns:
74 if not overwrite and os.path.exists(fn):
75 raise FileSaveError('file exists: %s' % fn)
77 ensuredirs(fn)
79 open_files[fn] = open(fn, ['wb', 'ab'][fn in fns])
80 writer = ims.Writer(open_files[fn])
81 writer.write(
82 ims.MessageHeader(
83 version='GSE2.1',
84 type='DATA',
85 msg_id=ims.MsgID(
86 msg_id_string=randomid(),
87 msg_id_source='Pyrocko_%s' % info.version)))
89 writer.write(ims.WaveformSection(
90 datatype=ims.DataType(
91 type='WAVEFORM',
92 format='GSE2.1')))
94 fns.add(fn)
96 sec = ims.WID2Section.from_pyrocko_trace(tr, None, None, None, None)
97 writer = ims.Writer(open_files[fn])
98 writer.write(sec)
100 for fn in fns:
101 if fn not in open_files:
102 open_files[fn] = open(fn, 'ab')
104 writer = ims.Writer(open_files[fn])
105 writer.write(ims.Stop())
106 open_files.pop(fn).close()
108 return list(fns)