1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import os
7import random
9from .io_common import FileLoadError, FileSaveError
10from pyrocko.util import ensuredirs
13def detect(first512):
14 lines = first512.lstrip().splitlines()
15 if len(lines) >= 2:
16 if lines[0].startswith(b'WID2 '):
17 return True
19 if lines[0].startswith(b'BEGIN GSE2'):
20 return True
22 if lines[0].startswith(b'DATA_TYPE WAVEFORM GSE2'):
23 return True
25 return False
28def iload(filename, load_data=True):
30 from . import ims
32 try:
33 with open(filename, 'rb') as f:
35 r = ims.Reader(f, load_data=load_data, version='GSE2.0',
36 dialect=None)
38 for sec in r:
39 if isinstance(sec, ims.WID2Section):
40 tr = sec.pyrocko_trace(checksum_error='warn')
41 yield tr
43 except (OSError, ims.DeserializeError) as e:
44 fle = FileLoadError(e)
45 fle.set_context('filename', filename)
46 raise fle
49def randomid():
50 return ''.join(chr(random.randint(97, 122)) for _ in range(20))
53def save(traces, filename_template, additional={}, max_open_files=10,
54 overwrite=True):
56 from pyrocko import info
57 from . import ims
59 fns = set()
60 open_files = {}
62 def close_files():
63 while open_files:
64 open_files.popitem()[1].close()
66 for tr in traces:
67 fn = tr.fill_template(filename_template, **additional)
68 if fn not in open_files:
69 if len(open_files) >= max_open_files:
70 close_files()
72 if fn not in fns:
73 if not overwrite and os.path.exists(fn):
74 raise FileSaveError('file exists: %s' % fn)
76 ensuredirs(fn)
78 open_files[fn] = open(fn, ['wb', 'ab'][fn in fns])
79 writer = ims.Writer(open_files[fn])
80 writer.write(
81 ims.MessageHeader(
82 version='GSE2.1',
83 type='DATA',
84 msg_id=ims.MsgID(
85 msg_id_string=randomid(),
86 msg_id_source='Pyrocko_%s' % info.version)))
88 writer.write(ims.WaveformSection(
89 datatype=ims.DataType(
90 type='WAVEFORM',
91 format='GSE2.1')))
93 fns.add(fn)
95 sec = ims.WID2Section.from_pyrocko_trace(tr, None, None, None, None)
96 writer = ims.Writer(open_files[fn])
97 writer.write(sec)
99 for fn in fns:
100 if fn not in open_files:
101 open_files[fn] = open(fn, 'ab')
103 writer = ims.Writer(open_files[fn])
104 writer.write(ims.Stop())
105 open_files.pop(fn).close()
107 return list(fns)