1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6from __future__ import absolute_import, print_function 

7 

8import copy 

9 

10try: 

11 from yaml import CSafeLoader as SafeLoader, CSafeDumper as SafeDumper 

12except ImportError: 

13 from yaml import SafeLoader, SafeDumper 

14 

15from pyrocko import guts 

16 

17 

18class AgnosticSafeLoader(SafeLoader): 

19 pass 

20 

21 

22class AgnosticSafeDumper(SafeDumper): 

23 pass 

24 

25 

26class xstr(str): 

27 pass 

28 

29 

30def to_xstr(x): 

31 if isinstance(x, str): 

32 return xstr(x) 

33 elif isinstance(x, list): 

34 return [to_xstr(e) for e in x] 

35 elif isinstance(x, dict): 

36 return dict((k, to_xstr(v)) for (k, v) in x.items()) 

37 else: 

38 return x 

39 

40 

41def quoted_presenter(dumper, data): 

42 return dumper.represent_scalar( 

43 'tag:yaml.org,2002:str', str(data), style='\'') 

44 

45 

46AgnosticSafeDumper.add_representer(xstr, quoted_presenter) 

47 

48 

49class Object(object): 

50 

51 def __init__(self, tagname, inamevals): 

52 self._tagname = tagname 

53 self._data = [] 

54 for kv in inamevals: 

55 self._data.append(list(kv)) 

56 

57 def inamevals_to_save(self): 

58 for k, v in self._data: 

59 yield (k, to_xstr(v)) 

60 

61 def inamevals(self): 

62 for k, v in self._data: 

63 yield (k, v) 

64 

65 def __iter__(self): 

66 for k, _ in self._data: 

67 yield k 

68 

69 def rename_attribute(self, old, new): 

70 for kv in self._data: 

71 if kv[0] == old: 

72 kv[0] = new 

73 

74 def drop_attribute(self, k): 

75 self._data = [kv for kv in self._data if kv[0] != k] 

76 

77 def replace(self, other): 

78 self._tagname = other._tagname 

79 self._data = copy.deepcopy(other._data) 

80 

81 def __setitem__(self, k, v): 

82 for kv in self._data: 

83 if kv[0] == k: 

84 kv[1] = v 

85 return 

86 

87 self._data.append([k, v]) 

88 

89 def __getitem__(self, item): 

90 for kv in self._data: 

91 if kv[0] == item: 

92 return kv[1] 

93 

94 raise KeyError(item) 

95 

96 def get(self, *args): 

97 if len(args) == 1: 

98 return self.__getitem__(args[0]) 

99 else: 

100 try: 

101 return self.__getitem__(args[0]) 

102 except KeyError: 

103 return args[1] 

104 

105 

106def multi_representer(dumper, data): 

107 node = dumper.represent_mapping( 

108 '!'+data._tagname, data.inamevals_to_save(), flow_style=False) 

109 

110 return node 

111 

112 

113def multi_constructor(loader, tag_suffix, node): 

114 tagname = str(tag_suffix) 

115 

116 o = Object(tagname, loader.construct_pairs(node, deep=True)) 

117 return o 

118 

119 

120AgnosticSafeDumper.add_multi_representer(Object, multi_representer) 

121AgnosticSafeLoader.add_multi_constructor('!', multi_constructor) 

122 

123 

124@guts.expand_stream_args('w') 

125def dump(*args, **kwargs): 

126 return guts._dump(Dumper=AgnosticSafeDumper, *args, **kwargs) 

127 

128 

129@guts.expand_stream_args('r') 

130def load(*args, **kwargs): 

131 return guts._load(Loader=AgnosticSafeLoader, *args, **kwargs) 

132 

133 

134def load_string(s, *args, **kwargs): 

135 return load(string=s, *args, **kwargs) 

136 

137 

138@guts.expand_stream_args('w') 

139def dump_all(*args, **kwargs): 

140 return guts._dump_all(Dumper=AgnosticSafeDumper, *args, **kwargs) 

141 

142 

143@guts.expand_stream_args('r') 

144def load_all(*args, **kwargs): 

145 return guts._load_all(Loader=AgnosticSafeLoader, *args, **kwargs) 

146 

147 

148@guts.expand_stream_args('r') 

149def iload_all(*args, **kwargs): 

150 return guts._iload_all(Loader=AgnosticSafeLoader, *args, **kwargs) 

151 

152 

153def walk(x, path=()): 

154 yield path, x 

155 

156 if isinstance(x, Object): 

157 for (name, val) in x.inamevals(): 

158 if isinstance(val, (list, tuple)): 

159 for iele, ele in enumerate(val): 

160 for y in walk(ele, path=path + ((name, iele),)): 

161 yield y 

162 elif isinstance(val, dict): 

163 for ele_k, ele_v in val.items(): 

164 for y in walk(ele_v, path=path + ((name, ele_k),)): 

165 yield y 

166 else: 

167 for y in walk(val, path=path+(name,)): 

168 yield y 

169 

170 

171def apply_tree(x, func, path=()): 

172 if isinstance(x, Object): 

173 for (name, val) in x.inamevals(): 

174 if isinstance(val, (list, tuple)): 

175 for iele, ele in enumerate(val): 

176 apply_tree(ele, func, path=path + ((name, iele),)) 

177 elif isinstance(val, dict): 

178 for ele_k, ele_v in val.items(): 

179 apply_tree(ele_v, func, path=path + ((name, ele_k),)) 

180 else: 

181 apply_tree(val, func, path=path+(name,)) 

182 

183 func(path, x)