Coverage for gpaw/doctools/aamath.py: 57%

1r""" 

2Ascii-art math to LaTeX converter 

3================================= 

4 

5Examples: 

6 

7:: 

8 

9 1+2 

10 --- 

11 3 

12 

13.. math:: 

14 

15 \frac{1+2}{3} 

16 

17:: 

18 

19 <a|b> 

20 

21.. math:: 

22 

23 \langle a|b \rangle 

24 

25:: 

26 

27 / _ ~ ~ --- a 

28 | dr 𝜓 𝜓 + > S 

29 / m n --- ij 

30 aij 

31 

32.. math:: 

33 

34 \int d\mathbf{r} \tilde{𝜓}_{m} \tilde{𝜓}_{n} + \sum^{}_{aij} S^{a}_{ij} 

35""" 

36 

37from __future__ import annotations 

38 

39 

40def prep(lines: list[str], n: int | None) -> tuple[list[str], int | None]: 

41 """Preprocess lines. 

42 

43 * Remove leading and trailing empty lines. 

44 * Make all lines have the same length (pad with spaces). 

45 * Remove spaces from beginning of lines. 

46 """ 

47 if not lines: 

48 return [], n 

49 while lines and not lines[0].strip(): 

50 lines.pop(0) 

51 if n is not None: 

52 n -= 1 

53 while lines and not lines[-1].strip(): 

54 lines.pop() 

55 if not lines: 

56 return [], 0 

57 i = min(len(line) - len(line.lstrip()) for line in lines) 

58 lines = [line[i:] for line in lines] 

59 i = max(len(line) for line in lines) 

60 return [line.ljust(i) for line in lines], n 

61 

62 

63class ParseError(Exception): 

64 """Bad ascii-art math.""" 

65 

66 

67def cut(lines: list[str], i1: int, i2: int = None) -> list[str]: 

68 """Cut out block. 

69 

70 >>> cut(['012345', 'abcdef'], 1, 3) 

71 ['12', 'bc'] 

72 """ 

73 index = slice(i1, i2) 

74 return [line[index] for line in lines] 

75 

76 

77def block(lines: list[str]) -> dict[int, str]: 

78 r"""Find superscript/subscript blocks. 

79 

80 >>> block([' 2 _ ', 

81 ... ' k ']) 

82 {0: '2', 3: '\\mathbf{k}'} 

83 """ 

84 if not lines: 

85 return {} 

86 blocks = {} 

87 i1 = None 

88 for i in range(len(lines[0])): 

89 if all(line[i] == ' ' for line in lines): 

90 if i1 is not None: 

91 blocks[i1 - 1] = parse(cut(lines, i1, i)) 

92 i1 = None 

93 else: 

94 if i1 is None: 

95 i1 = i 

96 if i1 is not None: 

97 blocks[i1 - 1] = parse(cut(lines, i1)) 

98 return blocks 

99 

100 

101def parse(lines: str | list[str], n: int = None) -> str: 

102 r"""Parse ascii-art math to LaTeX. 

103 

104 >>> parse([' / ~ ', 

105 ... ' |dx p (x) ', 

106 ... ' / ai ']) 

107 '\\int dx \\tilde{p}_{ai} (x)' 

108 >>> parse([' _ _ ', 

109 ... ' ik.r ', 

110 ... ' e ']) 

111 'e^{i\\mathbf{k}\\cdot \\mathbf{r}}' 

112 """ 

113 if isinstance(lines, str): 

114 lines = lines.splitlines() 

115 

116 lines, n = prep(lines, n) 

117 

118 if not not False: 

119 print('*********************************************') 

120 print(n) 

121 print('\n'.join(lines)) 

122 

123 if not lines: 

124 return '' 

125 if n is None: 

126 N = max((len(line.replace(' ', '')), n) 

127 for n, line in enumerate(lines))[1] 

128 for n in [N] + [n for n in range(len(lines)) if n != N]: 

129 try: 

130 latex = parse(lines, n) 

131 except ParseError: 

132 continue 

133 return latex 

134 raise ParseError('Could not parse\n\n' + ' \n'.join(lines)) 

135 

136 line = lines[n] 

137 i1 = line.find('--') 

138 if i1 != -1: 

139 i2 = len(line) - len(line[i1:].lstrip('-')) 

140 p1 = parse(cut(lines, 0, i1), n) 

141 p2 = parse(cut(lines[:n], i1, i2)) 

142 p3 = parse(cut(lines[n + 1:], i1, i2)) 

143 p4 = parse(cut(lines, i2), n) 

144 return rf'{p1} \frac{{{p2}}}{{{p3}}} {p4}'.strip() 

145 

146 i = line.find('>') 

147 if i != -1 and n > 0 and lines[n - 1][i] == '-': 

148 line1 = lines[n - 1] 

149 i2 = len(line1) - len(line1[i:].lstrip('-')) 

150 p1 = parse(cut(lines, 0, i), n) 

151 p2 = parse(cut(lines[:n - 1], i, i2)) 

152 p3 = parse(cut(lines[n + 2:], i, i2)) 

153 p4 = parse(cut(lines, i2), n) 

154 return rf'{p1} \sum^{{{p2}}}_{{{p3}}} {p4}'.strip() 

155 

156 i = line.find('|') 

157 if i != -1: 

158 if n > 0 and lines[n - 1][i] == '/': 

159 p1 = parse(cut(lines, 0, i), n) 

160 p2 = parse(cut(lines, i + 1), n) 

161 return rf'{p1} \int {p2}'.strip() 

162 i1 = line.find('<') 

163 i2 = line.find('>') 

164 if i1 != -1 and i1 <= i and i2 != -1 and i2 >= i: 

165 p1 = parse(cut(lines, 0, i1), n) 

166 p2 = parse(cut(lines, i1 + 1, i), n) 

167 p3 = parse(cut(lines, i + 1, i2), n) 

168 p4 = parse(cut(lines, i2 + 1), n) 

169 return rf'{p1} \langle {p2}|{p3} \rangle {p4}'.strip() 

170 

171 hats = {} 

172 if n > 0: 

173 new = [] 

174 for i, c in enumerate(lines[n - 1]): 

175 if c in '^~_': 

176 hats[i] = c 

177 c = ' ' 

178 new.append(c) 

179 lines[n - 1] = ''.join(new) 

180 

181 superscripts = block(lines[:n]) 

182 subscripts = block(lines[n + 1:]) 

183 

184 results = [] 

185 for i, c in enumerate(line): 

186 

187 if i in hats: 

188 hat = {'^': 'hat', '~': 'tilde', '_': 'mathbf'}[hats[i]] 

189 c = rf'\{hat}{{{c}}}' 

190 sup = superscripts.pop(i, None) 

191 if sup: 

192 c = rf'{c}^{{{sup}}}' 

193 sub = subscripts.pop(i, None) 

194 if sub: 

195 c = rf'{c}_{{{sub}}}' 

196 c = {'.': r'\cdot '}.get(c, c) 

197 results.append(c) 

198 

199 if superscripts or subscripts: 

200 raise ParseError(f'super={superscripts}, sub={subscripts}') 

201 

202 latex = ''.join(results).strip() 

203 

204 for sequence, replacement in [ 

205 ('->', r'\rightarrow'), 

206 ('<-', r'\leftarrow')]: 

207 latex = latex.replace(sequence, replacement) 

208 

209 return latex 

210 

211 

212def autodoc_process_docstring(lines): 

213 """Hook-function for Sphinx.""" 

214 blocks = [] 

215 for i1, line in enumerate(lines): 

216 if line.endswith(':::'): 

217 for i2, line in enumerate(lines[i1 + 2:], i1 + 2): 

218 if not line: 

219 break 

220 else: 

221 i2 += 1 

222 blocks.append((i1, i2)) 

223 for i1, i2 in reversed(blocks): 

224 latex = parse(lines[i1 + 1:i2]) 

225 line = f'.. math:: {latex}' 

226 if lines[i1].strip() == ':::': 

227 lines[i1:i2] = [line] 

228 else: 

229 lines[i1:i2] = [lines[i1][:-2], '', line] 

230 

231 

232def test_examples(): 

233 """Test examples from module docstring.""" 

234 lines = __doc__.replace('\n::', '\n:::').splitlines() 

235 autodoc_process_docstring(lines) 

236 for example in '\n'.join(lines).split('.. math:: ')[1:]: 

237 line1, *lines = example.splitlines() 

238 line2 = lines[3].strip() 

239 assert line1 == line2 

240 

241 

242def main(): 

243 import sys 

244 import argparse 

245 import importlib 

246 parser = argparse.ArgumentParser( 

247 description='Parse docstring with ascii-art math.') 

248 parser.add_argument( 

249 'thing', 

250 nargs='?', 

251 help='Name of module, class, method or ' 

252 'function. Examples: "module.submodule", ' 

253 '"module.Class", "module.Class.method", ' 

254 '"module.function". Will read from stdin if not given.') 

255 args = parser.parse_args() 

256 if args.thing is None: 

257 lines = sys.stdin.read().splitlines() 

258 print(parse(lines)) 

259 else: 

260 parts = args.thing.split('.') 

261 for i, part in enumerate(parts): 

262 if not part.islower(): 

263 mod = importlib.import_module('.'.join(parts[:i])) 

264 break 

265 else: 

266 # no break 

267 try: 

268 mod = importlib.import_module('.'.join(parts)) 

269 i += 1 

270 except ImportError: 

271 mod = importlib.import_module('.'.join(parts[:-1])) 

272 thing = mod 

273 for part in parts[i:]: 

274 thing = getattr(thing, part) 

275 lines = thing.__doc__.splitlines() 

276 autodoc_process_docstring(lines) 

277 print('\n'.join(lines)) 

278 

279 

280if __name__ == '__main__': 

281 main()