Coverage for /home/tbone/mambaforge/lib/python3.9/site-packages/importnb/loader.py: 67%

1# coding: utf-8

2"""# `loader`

3

4 Combine the __import__ finder with the loader.

5"""

6

7

8 import ast

9 import re

10 import sys

11 import textwrap

12 from dataclasses import asdict, dataclass, field

13 from functools import partial

14 from importlib import reload

15 from importlib._bootstrap import _init_module_attrs, _requires_builtin

16 from importlib._bootstrap_external import FileFinder, decode_source

17 from importlib.machinery import ModuleSpec, SourceFileLoader

18 from importlib.util import LazyLoader, find_spec

19 from pathlib import Path

20 from types import ModuleType

21

22 from . import get_ipython

23 from .decoder import LineCacheNotebookDecoder, quote

24 from .docstrings import update_docstring

25 from .finder import FuzzyFinder, get_loader_details, get_loader_index

26

27 _GTE38 = sys.version_info.major == 3 and sys.version_info.minor >= 8

28

29 try:

30 import IPython

31 from IPython.core.inputsplitter import IPythonInputSplitter

32

33 dedent = IPythonInputSplitter(

34 line_input_checker=False,

35 physical_line_transforms=[

36 IPython.core.inputsplitter.leading_indent(),

37 IPython.core.inputsplitter.ipy_prompt(),

38 IPython.core.inputsplitter.cellmagic(end_on_blank_line=False),

39 ],

40 ).transform_cell

41 except ModuleNotFoundError:

42

43 def dedent(body):

44 from textwrap import dedent, indent

45

46 if MAGIC.match(body):

47 return indent(body, "# ")

48 return dedent(body)

49

50

51 __all__ = "Notebook", "reload"

52

53

54 MAGIC = re.compile("^\s*%{2}", re.MULTILINE)

55

56

57@dataclass

58 class Interface:

59 name: str = None

60 path: str = None

61 lazy: bool = False

62 include_fuzzy_finder: bool = True

63

64 include_markdown_docstring: bool = True

65 only_defs: bool = False

66 no_magic: bool = False

67 _loader_hook_position: int = field(default=0, repr=False)

68

69 def __new__(cls, name=None, path=None, **kwargs):

70 kwargs.update(name=name, path=path)

71 self = super().__new__(cls)

72 self.__init__(**kwargs)

73 return self

74

75

76 class BaseLoader(Interface, SourceFileLoader):

77 """The simplest implementation of a Notebook Source File Loader."""

78

79 @property

80 def loader(self):

81 """Create a lazy loader source file loader."""

82 loader = type(self)

83 if self.lazy and (sys.version_info.major, sys.version_info.minor) != (3, 4):

84 loader = LazyLoader.factory(loader)

85 # Strip the leading underscore from slots

86 params = asdict(self)

87 params.pop("name")

88 params.pop("path")

89 return partial(loader, **params)

90

91 @property

92 def finder(self):

93 """Permit fuzzy finding of files with special characters."""

94 return self.include_fuzzy_finder and FuzzyFinder or FileFinder

95

96 def translate(self, source):

97 if self.path and self.path.endswith(".ipynb"):

98 return LineCacheNotebookDecoder(

99 code=self.code, raw=self.raw, markdown=self.markdown

100 ).decode(source, self.path)

101 return self.code(source)

102

103 def get_data(self, path):

104 """Needs to return the string source for the module."""

105 return self.translate(self.decode())

106

107 def create_module(self, spec):

108 module = ModuleType(str(spec.name))

109 _init_module_attrs(spec, module)

110 if self.name:

111 module.__name__ = self.name

112 if getattr(spec, "alias", None):

113 # put a fuzzy spec on the modules to avoid re importing it.

114 # there is a funky trick you do with the fuzzy finder where you

115 # load multiple versions with different finders.

116

117 sys.modules[spec.alias] = module

118 module.get_ipython = get_ipython

119 return module

120

121 def decode(self):

122 return decode_source(super().get_data(self.path))

123

124 def code(self, str):

125 return dedent(str)

126

127 def markdown(self, str):

128 return quote(str)

129

130 def raw(self, str):

131 return comment(str)

132

133 def visit(self, node):

134 return node

135

136 @classmethod

137 @_requires_builtin

138 def is_package(cls, fullname):

139 """Return False as built-in modules are never packages."""

140 if "." not in fullname:

141 return True

142 return super().is_package(fullname)

143

144 get_source = get_data

145

146 def __enter__(self):

147 path_id, loader_id, details = get_loader_index(".py")

148 for _, e in details:

149 if all(map(e.__contains__, self.extensions)):

150 self._loader_hook_position = None

151 return self

152 else:

153 self._loader_hook_position = loader_id + 1

154 details.insert(self._loader_hook_position, (self.loader, self.extensions))

155 sys.path_hooks[path_id] = self.finder.path_hook(*details)

156 sys.path_importer_cache.clear()

157 return self

158

159 def __exit__(self, *excepts):

160 if self._loader_hook_position is not None:

161 path_id, details = get_loader_details()

162 details.pop(self._loader_hook_position)

163 sys.path_hooks[path_id] = self.finder.path_hook(*details)

164 sys.path_importer_cache.clear()

165

166

167 class FileModuleSpec(ModuleSpec):

168 def __init__(self, *args, **kwargs):

169 super().__init__(*args, **kwargs)

170 self._set_fileattr = True

171

172

173 def comment(str):

174 return textwrap.indent(str, "# ")

175

176

177 class DefsOnly(ast.NodeTransformer):

178 INCLUDE = ast.Import, ast.ImportFrom, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef

179

180 def visit_Module(self, node):

181 args = ([x for x in node.body if isinstance(x, self.INCLUDE)],)

182 if _GTE38:

183 args += (node.type_ignores,)

184 return ast.Module(*args)

185

186

187 class Notebook(BaseLoader):

188 """Notebook is a user friendly file finder and module loader for notebook source code.

189

190 > Remember, restart and run all or it didn't happen.

191

192 Notebook provides several useful options.

193

194 * Lazy module loading. A module is executed the first time it is used in a script.

195 """

196

197 extensions = (".ipy", ".ipynb")

198

199 def parse(self, nodes):

200 return ast.parse(nodes, self.path)

201

202 def visit(self, nodes):

203 if self.only_defs:

204 nodes = DefsOnly().visit(nodes)

205 return nodes

206

207 def code(self, str):

208 if self.no_magic:

209 if MAGIC.match(str):

210 return comment(str)

211 return super().code(str)

212

213 def source_to_code(self, nodes, path, *, _optimize=-1):

214 """* Convert the current source to ast

215 * Apply ast transformers.

216 * Compile the code."""

217 if not isinstance(nodes, ast.Module):

218 nodes = self.parse(nodes)

219 if self.include_markdown_docstring:

220 nodes = update_docstring(nodes)

221 return super().source_to_code(

222 ast.fix_missing_locations(self.visit(nodes)), path, _optimize=_optimize

223 )

224

225 @classmethod

226 def load_file(cls, filename, main=True, **kwargs):

227 """Import a notebook as a module from a filename.

228

229 dir: The directory to load the file from.

230 main: Load the module in the __main__ context.

231

232 > assert Notebook.load('loader.ipynb')

233 """

234 name = main and "__main__" or filename

235 loader = cls(name, str(filename), **kwargs)

236 spec = FileModuleSpec(name, loader, origin=loader.path)

237 module = loader.create_module(spec)

238 loader.exec_module(module)

239 return module

240

241 @classmethod

242 def load_module(cls, module, main=False, **kwargs):

243 """Import a notebook as a module.

244

245 dir: The directory to load the file from.

246 main: Load the module in the __main__ context.

247

248 > assert Notebook.load('loader.ipynb')

249 """

250 from runpy import _run_module_as_main, run_module

251

252 with cls() as loader:

253 if main:

254 return _dict_module(_run_module_as_main(module))

255 else:

256 spec = find_spec(module)

257

258 m = spec.loader.create_module(spec)

259 spec.loader.exec_module(m)

260 return m

261

262 @classmethod

263 def load_argv(cls, argv=None, *, parser=None):

264 import sys

265 from sys import path

266

267 if parser is None:

268 parser = cls.get_argparser()

269

270 if argv is None:

271 from sys import argv

272

273 argv = argv[1:]

274

275 if isinstance(argv, str):

276 from shlex import split

277

278 argv = split(argv)

279

280 ns, unknown = parser.parse_known_args(argv)

281

282 if ns.code:

283 return cls.load_code(" ".join(ns.args))

284

285 n = ns.file or ns.module or sys.argv[0]

286

287 sys.argv = [n] + unknown

288 if ns.module:

289 path.insert(0, ns.dir) if ns.dir else ... if "" in path else path.insert(0, "")

290 return cls.load_module(ns.module, main=True)

291 elif ns.file:

292 if ns.dir:

293 n = str(Path(ns.dir) / ns.file)

294 return cls.load_file(n)

295 elif ns.code:

296 return cls.load_code(ns.code)

297 else:

298 parser.print_help()

299

300 @classmethod

301 def load_code(cls, code, mod_name=None, script_name=None, main=False):

302 from runpy import _run_module_code

303

304 self = cls()

305 name = main and "__main__" or mod_name or "<markdown code>"

306

307 return _dict_module(

308 _run_module_code(self.translate(code), mod_name=name, script_name=script_name)

309 )

310

311 @staticmethod

312 def get_argparser(parser=None):

313 from argparse import REMAINDER, ArgumentParser

314

315 if parser is None:

316 parser = ArgumentParser("importnb", description="run notebooks as python code")

317 parser.add_argument("-f", "--file", help="load a file")

318 parser.add_argument("-m", "--module", help="run args as a module")

319 parser.add_argument("-c", "--code", action="store_true", help="run args as code")

320 parser.add_argument("-d", "--dir", help="the directory path to run in.")

321 return parser

322

323

324 def _dict_module(ns):

325 m = ModuleType(ns.get("__name__"), ns.get("__doc__"))

326 m.__dict__.update(ns)

327 return m