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

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

205 statements  

1# coding: utf-8 

2"""# `loader` 

3 

4Combine the __import__ finder with the loader. 

5""" 

6 

7 

8import ast 

9import re 

10import sys 

11import textwrap 

12from dataclasses import asdict, dataclass, field 

13from functools import partial 

14from importlib import reload 

15from importlib._bootstrap import _init_module_attrs, _requires_builtin 

16from importlib._bootstrap_external import FileFinder, decode_source 

17from importlib.machinery import ModuleSpec, SourceFileLoader 

18from importlib.util import LazyLoader, find_spec 

19from pathlib import Path 

20from types import ModuleType 

21 

22from . import get_ipython 

23from .decoder import LineCacheNotebookDecoder, quote 

24from .docstrings import update_docstring 

25from .finder import FuzzyFinder, get_loader_details, get_loader_index 

26 

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

28 

29try: 

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 

41except 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 

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

55 

56 

57@dataclass 

58class 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 

76class 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 

167class FileModuleSpec(ModuleSpec): 

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

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

170 self._set_fileattr = True 

171 

172 

173def comment(str): 

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

175 

176 

177class 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 

187class 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 

324def _dict_module(ns): 

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

326 m.__dict__.update(ns) 

327 return m