|
# http://pyrocko.org - GPLv3
#
# The Pyrocko Developers, 21st Century
# ---|P------/S----------~Lg----------
from __future__ import division, absolute_import
from builtins import range
from builtins import chr
import os
import random
from .io_common import FileLoadError, FileSaveError
from pyrocko.util import ensuredirs
def detect(first512):
lines = first512.lstrip().splitlines()
if len(lines) >= 2:
if lines[0].startswith(b'WID2 '):
return True
if lines[0].startswith(b'BEGIN GSE2'):
return True
if lines[0].startswith(b'DATA_TYPE WAVEFORM GSE2'):
return True
return False
def iload(filename, load_data=True):
from . import ims
try:
with open(filename, 'rb') as f:
r = ims.Reader(f, load_data=load_data, version='GSE2.0',
dialect=None)
for sec in r:
if isinstance(sec, ims.WID2Section):
tr = sec.pyrocko_trace(checksum_error='warn')
yield tr
except (OSError, ims.DeserializeError) as e:
fle = FileLoadError(e)
fle.set_context('filename', filename)
raise fle
def randomid():
return ''.join(chr(random.randint(97, 122)) for _ in range(20))
def save(traces, filename_template, additional={}, max_open_files=10,
overwrite=True):
from pyrocko import info
from . import ims
fns = set()
open_files = {}
def close_files():
while open_files:
open_files.popitem()[1].close()
for tr in traces:
fn = tr.fill_template(filename_template, **additional)
if fn not in open_files:
if len(open_files) >= max_open_files:
close_files()
if fn not in fns:
if not overwrite and os.path.exists(fn):
raise FileSaveError('file exists: %s' % fn)
ensuredirs(fn)
open_files[fn] = open(fn, ['wb', 'ab'][fn in fns])
writer = ims.Writer(open_files[fn])
writer.write(
ims.MessageHeader(
version='GSE2.1',
type='DATA',
msg_id=ims.MsgID(
msg_id_string=randomid(),
msg_id_source='Pyrocko_%s' % info.version)))
writer.write(ims.WaveformSection(
datatype=ims.DataType(
type='WAVEFORM',
format='GSE2.1')))
fns.add(fn)
sec = ims.WID2Section.from_pyrocko_trace(tr, None, None, None, None)
writer = ims.Writer(open_files[fn])
writer.write(sec)
for fn in fns:
if fn not in open_files:
open_files[fn] = open(fn, 'ab')
writer = ims.Writer(open_files[fn])
writer.write(ims.Stop())
open_files.pop(fn).close()
return list(fns)
|