Coverage for pyquickhelper/ipythonhelper/notebook_helper.py: 67%

163 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-03 02:21 +0200

1""" 

2@file 

3@brief Some automation helpers about notebooks 

4""" 

5import os 

6import sys 

7import json 

8import warnings 

9from io import StringIO 

10from nbformat import versions 

11from nbformat.reader import reads, NotJSONError 

12from nbformat.v4 import upgrade 

13from ..filehelper import read_content_ufs 

14from ..loghelper import noLOG 

15from ..filehelper import explore_folder_iterfile, remove_folder 

16from .notebook_runner import NotebookRunner 

17from .notebook_exception import NotebookException 

18 

19 

20with warnings.catch_warnings(): 

21 warnings.simplefilter("ignore", category=ImportWarning) 

22 try: 

23 from ipykernel.kernelspec import install as install_k 

24 raisewarn = False 

25 except ImportError: # pragma: no cover 

26 raisewarn = True 

27if raisewarn: # pragma: no cover 

28 warnings.warn("ipykernel is not installed. pyquickhelper cannot execute a notebook.", 

29 category=ImportWarning) 

30 

31 

32def writes(nb, **kwargs): 

33 """ 

34 Write a notebook to a string in a given format in the current nbformat version. 

35 

36 This function always writes the notebook in the current nbformat version. 

37 

38 Parameters 

39 ++++++++++ 

40 

41 nb : NotebookNode 

42 The notebook to write. 

43 kwargs : 

44 Among these parameters, *version* (int) which is 

45 The nbformat version to write. 

46 Used for downgrading notebooks. 

47 

48 Returns 

49 +++++++ 

50 

51 s : unicode 

52 The notebook string. 

53 """ 

54 try: 

55 return versions[nb.nbformat].writes_json(nb, **kwargs) 

56 except AttributeError as e: # pragma: no cover 

57 raise NotebookException( 

58 f"probably wrong error: {nb.nbformat}") from e 

59 

60 

61def upgrade_notebook(filename, encoding="utf-8"): 

62 """ 

63 Converts a notebook from version 2 to latest. 

64 

65 @param filename filename 

66 @param encoding encoding 

67 @return modification? 

68 """ 

69 with open(filename, "r", encoding=encoding) as payload: 

70 content = payload.read() 

71 

72 try: 

73 nb = reads(content) 

74 except NotJSONError as e: # pragma: no cover 

75 if len(content) > 10: 

76 lc = list(content[:10]) 

77 else: 

78 lc = list(content) 

79 raise ValueError( 

80 f"Unable to read content type '{type(content)}' in '{filename}' ---- {lc}") from e 

81 

82 if not hasattr(nb, "nbformat") or nb.nbformat >= 4: 

83 return False 

84 

85 try: 

86 upgrade(nb, from_version=nb.nbformat) 

87 except ValueError as e: # pragma: no cover 

88 raise ValueError(f"Unable to convert '{filename}'.") from e 

89 

90 s = writes(nb) 

91 if isinstance(s, bytes): 

92 s = s.decode('utf8') 

93 

94 if s == content: 

95 return False 

96 with open(filename, "w", encoding=encoding) as f: 

97 f.write(s) 

98 return True 

99 

100 

101def read_nb(filename, profile_dir=None, encoding="utf8", working_dir=None, 

102 comment="", fLOG=noLOG, code_init=None, 

103 kernel_name="python", log_level="30", extended_args=None, 

104 kernel=False, replacements=None): 

105 """ 

106 Reads a notebook and return a @see cl NotebookRunner object. 

107 

108 @param filename notebook filename (or stream) 

109 @param profile_dir profile directory 

110 @param encoding encoding for the notebooks 

111 @param working_dir working directory 

112 @param comment additional information added to error message 

113 @param code_init to initialize the notebook with a python code as if it was a cell 

114 @param fLOG logging function 

115 @param log_level Choices: (0, 10, 20, 30=default, 40, 50, 'DEBUG', 'INFO', 'WARN', 'ERROR', 'CRITICAL') 

116 @param kernel_name kernel name, it can be None 

117 @param extended_args others arguments to pass to the command line 

118 (`--KernelManager.autorestar=True` for example), 

119 see :ref:`l-ipython_notebook_args` for a full list 

120 @param kernel *kernel* is True by default, the notebook can be run, if False, 

121 the notebook can be read but not run 

122 @param replacements replacements to make in every cell before running it, 

123 dictionary ``{ string: string }`` 

124 @return @see cl NotebookRunner 

125 """ 

126 if isinstance(filename, str): 

127 with open(filename, "r", encoding=encoding) as payload: 

128 nb = reads(payload.read()) 

129 

130 nb_runner = NotebookRunner( 

131 nb, profile_dir=profile_dir, theNotebook=os.path.abspath(filename), 

132 kernel=kernel, working_dir=working_dir, 

133 comment=comment, fLOG=fLOG, code_init=code_init, 

134 kernel_name="python", log_level="30", extended_args=None, 

135 filename=filename, replacements=replacements) 

136 return nb_runner 

137 else: 

138 nb = reads(filename.read()) 

139 nb_runner = NotebookRunner(nb, kernel=kernel, 

140 profile_dir=profile_dir, working_dir=working_dir, 

141 comment=comment, fLOG=fLOG, code_init=code_init, 

142 kernel_name="python", log_level="30", extended_args=None, 

143 filename=filename, replacements=replacements) 

144 return nb_runner 

145 

146 

147def read_nb_json(js, profile_dir=None, encoding="utf8", 

148 working_dir=None, comment="", fLOG=noLOG, code_init=None, 

149 kernel_name="python", log_level="30", extended_args=None, 

150 kernel=False, replacements=None): 

151 """ 

152 Reads a notebook from a :epkg:`JSON` stream or string. 

153 

154 @param js string or stream 

155 @param profile_dir profile directory 

156 @param encoding encoding for the notebooks 

157 @param working_dir working directory 

158 @param comment additional information added to error message 

159 @param code_init to initialize the notebook with a python code as if it was a cell 

160 @param fLOG logging function 

161 @param log_level Choices: (0, 10, 20, 30=default, 40, 50, 'DEBUG', 'INFO', 'WARN', 'ERROR', 'CRITICAL') 

162 @param kernel_name kernel name, it can be None 

163 @param extended_args others arguments to pass to the command line ('--KernelManager.autorestar=True' for example), 

164 see :ref:`l-ipython_notebook_args` for a full list 

165 @param kernel *kernel* is True by default, the notebook can be run, if False, 

166 the notebook can be read but not run 

167 @param replacements replacements to make in every cell before running it, 

168 dictionary ``{ string: string }`` 

169 @return instance of @see cl NotebookRunner 

170 """ 

171 if isinstance(js, str): 

172 st = StringIO(js) 

173 else: 

174 st = js 

175 return read_nb(st, encoding=encoding, kernel=kernel, 

176 profile_dir=profile_dir, working_dir=working_dir, 

177 comment=comment, fLOG=fLOG, code_init=code_init, 

178 kernel_name="python", log_level="30", extended_args=None, 

179 replacements=replacements) 

180 

181 

182def find_notebook_kernel(kernel_spec_manager=None): 

183 """ 

184 Returns a dict mapping kernel names to resource directories. 

185 

186 @param kernel_spec_manager see `KernelSpecManager <http://jupyter-client.readthedocs.org/en/ 

187 latest/api/kernelspec.html#jupyter_client.kernelspec.KernelSpecManager>`_ 

188 A KernelSpecManager to use for installation. 

189 If none provided, a default instance will be created. 

190 @return dict 

191 

192 The list of installed kernels is described at 

193 `Making kernel for Jupyter <http://jupyter-client.readthedocs.org/en/latest/kernels.html#kernelspecs>`_. 

194 The function only works with *Jupyter>=4.0*. 

195 """ 

196 if kernel_spec_manager is None: 

197 from jupyter_client.kernelspec import KernelSpecManager 

198 kernel_spec_manager = KernelSpecManager() 

199 return kernel_spec_manager.find_kernel_specs() 

200 

201 

202def get_notebook_kernel(kernel_name, kernel_spec_manager=None): 

203 """ 

204 Returns a `KernelSpec <https://ipython.org/ipython-doc/dev/api/ 

205 generated/IPython.kernel.kernelspec.html>`_. 

206 

207 @param kernel_spec_manager see `KernelSpecManager <http://jupyter-client.readthedocs.org/en/ 

208 latest/api/kernelspec.html#jupyter_client.kernelspec.KernelSpecManager>`_ 

209 A KernelSpecManager to use for installation. 

210 If none provided, a default instance will be created. 

211 @param kernel_name kernel name 

212 @return KernelSpec 

213 

214 The function only works with *Jupyter>=4.0*. 

215 """ 

216 if kernel_spec_manager is None: 

217 from jupyter_client.kernelspec import KernelSpecManager 

218 kernel_spec_manager = KernelSpecManager() 

219 return kernel_spec_manager.get_kernel_spec(kernel_name) 

220 

221 

222def install_notebook_extension(path=None, overwrite=False, symlink=False, 

223 user=False, prefix=None, nbextensions_dir=None, 

224 destination=None): 

225 """ 

226 Installs notebook extensions, 

227 see `install_nbextension <https://ipython.org/ipython-doc/ 

228 dev/api/generated/IPython.html.nbextensions.html 

229 #IPython.html.nbextensions.install_nbextension>`_ 

230 for documentation. 

231 

232 @param path if None, use default value 

233 @param overwrite overwrite the extension 

234 @param symlink see the original function 

235 @param user user 

236 @param prefix see the original function 

237 @param nbextensions_dir see the original function 

238 @param destination see the original function 

239 @return standard output 

240 

241 Default value is 

242 `https://github.com/ipython-contrib/IPython-notebook-extensions/archive/master.zip 

243 <https://github.com/ipython-contrib/IPython-notebook-extensions/archive/master.zip>`_. 

244 """ 

245 if path is None: 

246 path = "https://github.com/ipython-contrib/IPython-notebook-extensions/archive/master.zip" 

247 

248 cout = sys.stdout 

249 cerr = sys.stderr 

250 sys.stdout = StringIO() 

251 sys.stderr = StringIO() 

252 from notebook.nbextensions import install_nbextension 

253 install_nbextension(path=path, overwrite=overwrite, symlink=symlink, 

254 user=user, prefix=prefix, nbextensions_dir=nbextensions_dir, 

255 destination=destination) 

256 

257 out = sys.stdout.getvalue() 

258 err = sys.stderr.getvalue() 

259 sys.stdout = cout 

260 sys.stderr = cerr 

261 if len(err) != 0: 

262 raise NotebookException( 

263 f"unable to install exception from: {path}\nOUT:\n{out}\n[nberror]\n{err}") 

264 return out 

265 

266 

267def get_jupyter_datadir(): 

268 """ 

269 Returns the data directory for the notebook. 

270 

271 @return path 

272 """ 

273 from jupyter_client.kernelspec import KernelSpecManager 

274 return KernelSpecManager().data_dir 

275 

276 

277def get_jupyter_extension_dir(user=False, prefix=None, 

278 nbextensions_dir=None): 

279 """ 

280 Parameters 

281 ++++++++++ 

282 

283 user : bool [default: False] 

284 Whether to check the user's .ipython/nbextensions directory. 

285 Otherwise check a system-wide install (e.g. /usr/local/share/jupyter/nbextensions). 

286 prefix : str [optional] 

287 Specify install prefix, if it should differ from default (e.g. /usr/local). 

288 Will check prefix/share/jupyter/nbextensions 

289 nbextensions_dir : str [optional] 

290 Specify absolute path of nbextensions directory explicitly. 

291 

292 Return 

293 ++++++ 

294 

295 path: path to installed extensions (by the user) 

296 """ 

297 from notebook.nbextensions import _get_nbextension_dir 

298 return _get_nbextension_dir(nbextensions_dir=nbextensions_dir, user=user, prefix=prefix) 

299 

300 

301def get_installed_notebook_extension(user=False, prefix=None, 

302 nbextensions_dir=None): 

303 """ 

304 Retuns installed extensions. 

305 

306 :param user: bool [default: False] 

307 Whether to check the user's .ipython/nbextensions directory. 

308 Otherwise check a system-wide install (e.g. /usr/local/share/jupyter/nbextensions). 

309 :param prefix: str [optional] 

310 Specify install prefix, if it should differ from default (e.g. /usr/local). 

311 Will check prefix/share/jupyter/nbextensions 

312 :param nbextensions_dir: str [optional] 

313 Specify absolute path of nbextensions directory explicitly. 

314 :return: list: list of installed notebook extension (by the user) 

315 

316 You can install extensions with function @see fn install_notebook_extension. 

317 """ 

318 path = get_jupyter_extension_dir( 

319 user=user, prefix=prefix, nbextensions_dir=nbextensions_dir) 

320 if not os.path.exists(path): 

321 raise FileNotFoundError( # pragma: no cover 

322 f"Unable to find {path!r}.") 

323 

324 res = [] 

325 for file in explore_folder_iterfile(path): 

326 rel = os.path.relpath(file, path) 

327 spl = os.path.split(rel) 

328 name = spl[-1] 

329 if name == "main.js": 

330 fold = "/".join(spl[:-1]).replace("\\", "/") + "/main" 

331 res.append(fold) 

332 return res 

333 

334 

335def install_jupyter_kernel(exe=sys.executable, kernel_spec_manager=None, user=False, kernel_name=None, prefix=None): 

336 """ 

337 Installs a kernel based on executable (this python by default). 

338 

339 @param exe Python executable 

340 current one by default 

341 @param kernel_spec_manager (KernelSpecManager [optional]). 

342 A KernelSpecManager to use for installation. 

343 If none provided, a default instance will be created. 

344 @param user (bool). 

345 Whether to do a user-only install, or system-wide. 

346 @param kernel_name (str), optional. 

347 Specify a name for the kernelspec. 

348 This is needed for having multiple IPython 

349 kernels for different environments. 

350 @param prefix (str), optional. 

351 Specify an install prefix for the kernelspec. 

352 This is needed to install into a non-default 

353 location, such as a conda/virtual-env. 

354 

355 @return The path where the kernelspec was installed. 

356 

357 A kernel is defined by the following fields: 

358 

359 :: 

360 

361 { 

362 "display_name": "Python 3 (ENSAE)", 

363 "language": "python", 

364 "argv": [ "c:\\\\PythonENSAE\\\\python\\\\python.exe", 

365 "-m", 

366 "ipykernel", 

367 "-f", 

368 "{connection_file}" 

369 ] 

370 } 

371 

372 For R, it looks like: 

373 

374 :: 

375 

376 { 

377 "display_name": "R (ENSAE)", 

378 "language": "R", 

379 "argv": [ "c:\\\\PythonENSAE\\\\tools\\\\R\\\\bin\\\\x64\\\\R.exe", 

380 "--quiet", 

381 "-e", 

382 "IRkernel::main()", 

383 "--args", 

384 "{connection_file}" 

385 ] 

386 } 

387 """ 

388 exe = exe.replace("pythonw.exe", "python.exe") 

389 dest = install_k(kernel_spec_manager=kernel_spec_manager, 

390 user=user, kernel_name=kernel_name, prefix=prefix) 

391 kernel_file = os.path.join(dest, "kernel.json") 

392 kernel = dict(display_name=kernel_name, 

393 language="python", 

394 argv=[exe, "-m", "ipykernel", "-f", "{connection_file}"]) 

395 

396 s = json.dumps(kernel) 

397 with open(kernel_file, "w") as f: 

398 f.write(s) 

399 

400 return dest 

401 

402 

403def install_python_kernel_for_unittest(suffix=None): 

404 """ 

405 Installs a kernel based on this python (sys.executable) for unit test purposes. 

406 

407 @param suffix suffix to add to the kernel name 

408 @return kernel name 

409 """ 

410 exe = os.path.split(sys.executable)[0].replace("pythonw", "python") 

411 exe = exe.replace("\\", "/").replace("/", 

412 "_").replace(".", "_").replace(":", "") 

413 kern = "ut_" + exe + "_" + str(sys.version_info[0]) 

414 if suffix is not None: 

415 kern += "_" + suffix 

416 kern = kern.lower() 

417 install_jupyter_kernel(kernel_name=kern) 

418 return kern 

419 

420 

421def remove_kernel(kernel_name, kernel_spec_manager=None): 

422 """ 

423 Removes a kernel. 

424 

425 @param kernel_spec_manager see `KernelSpecManager <http://jupyter-client.readthedocs.org/ 

426 en/latest/api/kernelspec.html#jupyter_client.kernelspec.KernelSpecManager>`_ 

427 A KernelSpecManager to use for installation. 

428 If none provided, a default instance will be created. 

429 @param kernel_name kernel name 

430 

431 The function only works with *Jupyter>=4.0*. 

432 """ 

433 kernels = find_notebook_kernel(kernel_spec_manager=kernel_spec_manager) 

434 if kernel_name in kernels: 

435 fold = kernels[kernel_name] 

436 if not os.path.exists(fold): 

437 raise FileNotFoundError( # pragma: no cover 

438 f"Unable to remove folder {fold!r}.") 

439 remove_folder(fold) 

440 else: 

441 raise NotebookException( # pragma: no cover 

442 f"Unable to find kernel '{kernel_name}' in {', '.join(kernels.keys())}") 

443 

444 

445def remove_execution_number(infile, outfile=None, encoding="utf-8", indent=2, rule=int): 

446 """ 

447 Removes execution number from a notebook. 

448 

449 @param infile filename of the notebook 

450 @param outfile None ot save the file 

451 @param encoding encoding 

452 @param indent indentation 

453 @param rule determines the rule which specifies execution numbers, 

454 'None' for None, 'int' for consectuive integers numbers. 

455 @return modified string or None if outfile is not None 

456 and the file was not modified 

457 

458 .. todoext:: 

459 :title: remove execution number from notebook facilitate git versionning 

460 :tag: enhancement 

461 :issue: 18 

462 :cost: 1 

463 :hidden: 

464 :date: 2016-08-23 

465 :release: 1.4 

466 

467 Remove execution number from the notebook 

468 to avoid commiting changes only about those numbers 

469 

470 `notebook 5.1.0 

471 <https://jupyter-notebook.readthedocs.io/en/stable/changelog.html>`_ 

472 introduced changes which are incompatible with 

473 leaving the cell executing number empty. 

474 """ 

475 def fixup(adict, k, v, cellno=0, outputs="outputs"): 

476 for key in adict.keys(): 

477 if key == k: 

478 if rule is None: 

479 adict[key] = v 

480 elif rule is int: 

481 cellno += 1 

482 adict[key] = cellno 

483 else: 

484 raise ValueError( # pragma: no cover 

485 f"Rule '{rule}' does not apply on {key}={adict[key]}") 

486 elif key == "outputs": 

487 if isinstance(adict[key], dict): 

488 fixup(adict[key], k, v, cellno=cellno, outputs=outputs) 

489 elif isinstance(adict[key], list): 

490 for el in adict[key]: 

491 if isinstance(el, dict): 

492 fixup(el, k, v, cellno=cellno, outputs=outputs) 

493 elif isinstance(adict[key], dict): 

494 cellno = fixup(adict[key], k, v, 

495 cellno=cellno, outputs=outputs) 

496 elif isinstance(adict[key], list): 

497 for el in adict[key]: 

498 if isinstance(el, dict): 

499 cellno = fixup(el, k, v, cellno=cellno, 

500 outputs=outputs) 

501 return cellno 

502 

503 def widget(adict): 

504 metadata = adict.get('metadata', None) 

505 if metadata is None: 

506 return 

507 if 'widgets' in metadata: 

508 del metadata['widgets'] 

509 

510 content = read_content_ufs(infile) 

511 js = json.loads(content) 

512 fixup(js, "execution_count", None) 

513 widget(js) 

514 st = StringIO() 

515 json.dump(js, st, indent=indent, sort_keys=True) 

516 res = st.getvalue() 

517 if outfile is not None: 

518 if content != res: 

519 with open(outfile, "w", encoding=encoding) as f: 

520 f.write(res) 

521 return content 

522 return None 

523 return res