Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/io/gse2.py: 85%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6''' 

7Reader for GSE2.0 (IMS) waveform files. 

8''' 

9 

10import os 

11import random 

12 

13from .io_common import FileLoadError, FileSaveError 

14from pyrocko.util import ensuredirs 

15 

16 

17def detect(first512): 

18 lines = first512.lstrip().splitlines() 

19 if len(lines) >= 2: 

20 if lines[0].startswith(b'WID2 '): 

21 return True 

22 

23 if lines[0].startswith(b'BEGIN GSE2'): 

24 return True 

25 

26 if lines[0].startswith(b'DATA_TYPE WAVEFORM GSE2'): 

27 return True 

28 

29 return False 

30 

31 

32def iload(filename, load_data=True): 

33 

34 from . import ims 

35 

36 try: 

37 with open(filename, 'rb') as f: 

38 

39 r = ims.Reader(f, load_data=load_data, version='GSE2.0', 

40 dialect=None) 

41 

42 for sec in r: 

43 if isinstance(sec, ims.WID2Section): 

44 tr = sec.pyrocko_trace(checksum_error='warn') 

45 yield tr 

46 

47 except (OSError, ims.DeserializeError) as e: 

48 fle = FileLoadError(e) 

49 fle.set_context('filename', filename) 

50 raise fle 

51 

52 

53def randomid(): 

54 return ''.join(chr(random.randint(97, 122)) for _ in range(20)) 

55 

56 

57def save(traces, filename_template, additional={}, max_open_files=10, 

58 overwrite=True): 

59 

60 from pyrocko import info 

61 from . import ims 

62 

63 fns = set() 

64 open_files = {} 

65 

66 def close_files(): 

67 while open_files: 

68 open_files.popitem()[1].close() 

69 

70 for tr in traces: 

71 fn = tr.fill_template(filename_template, **additional) 

72 if fn not in open_files: 

73 if len(open_files) >= max_open_files: 

74 close_files() 

75 

76 if fn not in fns: 

77 if not overwrite and os.path.exists(fn): 

78 raise FileSaveError('file exists: %s' % fn) 

79 

80 ensuredirs(fn) 

81 

82 open_files[fn] = open(fn, ['wb', 'ab'][fn in fns]) 

83 writer = ims.Writer(open_files[fn]) 

84 writer.write( 

85 ims.MessageHeader( 

86 version='GSE2.1', 

87 type='DATA', 

88 msg_id=ims.MsgID( 

89 msg_id_string=randomid(), 

90 msg_id_source='Pyrocko_%s' % info.version))) 

91 

92 writer.write(ims.WaveformSection( 

93 datatype=ims.DataType( 

94 type='WAVEFORM', 

95 format='GSE2.1'))) 

96 

97 fns.add(fn) 

98 

99 sec = ims.WID2Section.from_pyrocko_trace(tr, None, None, None, None) 

100 writer = ims.Writer(open_files[fn]) 

101 writer.write(sec) 

102 

103 for fn in fns: 

104 if fn not in open_files: 

105 open_files[fn] = open(fn, 'ab') 

106 

107 writer = ims.Writer(open_files[fn]) 

108 writer.write(ims.Stop()) 

109 open_files.pop(fn).close() 

110 

111 return list(fns)