1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import copy 

7 

8try: 

9 from yaml import CSafeLoader as SafeLoader, CSafeDumper as SafeDumper 

10except ImportError: 

11 from yaml import SafeLoader, SafeDumper 

12 

13from pyrocko import guts 

14 

15 

16class AgnosticSafeLoader(SafeLoader): 

17 pass 

18 

19 

20class AgnosticSafeDumper(SafeDumper): 

21 pass 

22 

23 

24class xstr(str): 

25 pass 

26 

27 

28def to_xstr(x): 

29 if isinstance(x, str): 

30 return xstr(x) 

31 elif isinstance(x, list): 

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

33 elif isinstance(x, dict): 

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

35 else: 

36 return x 

37 

38 

39def quoted_presenter(dumper, data): 

40 return dumper.represent_scalar( 

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

42 

43 

44AgnosticSafeDumper.add_representer(xstr, quoted_presenter) 

45 

46 

47class Object(object): 

48 

49 def __init__(self, tagname, inamevals): 

50 self._tagname = tagname 

51 self._data = [] 

52 for kv in inamevals: 

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

54 

55 def inamevals_to_save(self): 

56 for k, v in self._data: 

57 yield (k, to_xstr(v)) 

58 

59 def inamevals(self): 

60 for k, v in self._data: 

61 yield (k, v) 

62 

63 def __iter__(self): 

64 for k, _ in self._data: 

65 yield k 

66 

67 def rename_attribute(self, old, new): 

68 for kv in self._data: 

69 if kv[0] == old: 

70 kv[0] = new 

71 

72 def drop_attribute(self, k): 

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

74 

75 def replace(self, other): 

76 self._tagname = other._tagname 

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

78 

79 def __setitem__(self, k, v): 

80 for kv in self._data: 

81 if kv[0] == k: 

82 kv[1] = v 

83 return 

84 

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

86 

87 def __getitem__(self, item): 

88 for kv in self._data: 

89 if kv[0] == item: 

90 return kv[1] 

91 

92 raise KeyError(item) 

93 

94 def get(self, *args): 

95 if len(args) == 1: 

96 return self.__getitem__(args[0]) 

97 else: 

98 try: 

99 return self.__getitem__(args[0]) 

100 except KeyError: 

101 return args[1] 

102 

103 

104def multi_representer(dumper, data): 

105 node = dumper.represent_mapping( 

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

107 

108 return node 

109 

110 

111def multi_constructor(loader, tag_suffix, node): 

112 tagname = str(tag_suffix) 

113 

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

115 return o 

116 

117 

118AgnosticSafeDumper.add_multi_representer(Object, multi_representer) 

119AgnosticSafeLoader.add_multi_constructor('!', multi_constructor) 

120 

121 

122@guts.expand_stream_args('w') 

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

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

125 

126 

127@guts.expand_stream_args('r') 

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

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

130 

131 

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

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

134 

135 

136@guts.expand_stream_args('w') 

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

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

139 

140 

141@guts.expand_stream_args('r') 

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

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

144 

145 

146@guts.expand_stream_args('r') 

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

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

149 

150 

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

152 yield path, x 

153 

154 if isinstance(x, Object): 

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

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

157 for iele, ele in enumerate(val): 

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

159 yield y 

160 elif isinstance(val, dict): 

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

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

163 yield y 

164 else: 

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

166 yield y 

167 

168 

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

170 if isinstance(x, Object): 

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

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

173 for iele, ele in enumerate(val): 

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

175 elif isinstance(val, dict): 

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

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

178 else: 

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

180 

181 func(path, x)