IpythonMagic.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565
  1. # -*- coding: utf-8 -*-
  2. """
  3. =====================
  4. Cython related magics
  5. =====================
  6. Magic command interface for interactive work with Cython
  7. .. note::
  8. The ``Cython`` package needs to be installed separately. It
  9. can be obtained using ``easy_install`` or ``pip``.
  10. Usage
  11. =====
  12. To enable the magics below, execute ``%load_ext cython``.
  13. ``%%cython``
  14. {CYTHON_DOC}
  15. ``%%cython_inline``
  16. {CYTHON_INLINE_DOC}
  17. ``%%cython_pyximport``
  18. {CYTHON_PYXIMPORT_DOC}
  19. Author:
  20. * Brian Granger
  21. Code moved from IPython and adapted by:
  22. * Martín Gaitán
  23. Parts of this code were taken from Cython.inline.
  24. """
  25. #-----------------------------------------------------------------------------
  26. # Copyright (C) 2010-2011, IPython Development Team.
  27. #
  28. # Distributed under the terms of the Modified BSD License.
  29. #
  30. # The full license is in the file ipython-COPYING.rst, distributed with this software.
  31. #-----------------------------------------------------------------------------
  32. from __future__ import absolute_import, print_function
  33. import imp
  34. import io
  35. import os
  36. import re
  37. import sys
  38. import time
  39. import copy
  40. import distutils.log
  41. import textwrap
  42. IO_ENCODING = sys.getfilesystemencoding()
  43. IS_PY2 = sys.version_info[0] < 3
  44. try:
  45. reload
  46. except NameError: # Python 3
  47. from imp import reload
  48. try:
  49. import hashlib
  50. except ImportError:
  51. import md5 as hashlib
  52. from distutils.core import Distribution, Extension
  53. from distutils.command.build_ext import build_ext
  54. from IPython.core import display
  55. from IPython.core import magic_arguments
  56. from IPython.core.magic import Magics, magics_class, cell_magic
  57. try:
  58. from IPython.paths import get_ipython_cache_dir
  59. except ImportError:
  60. # older IPython version
  61. from IPython.utils.path import get_ipython_cache_dir
  62. from IPython.utils.text import dedent
  63. from ..Shadow import __version__ as cython_version
  64. from ..Compiler.Errors import CompileError
  65. from .Inline import cython_inline
  66. from .Dependencies import cythonize
  67. PGO_CONFIG = {
  68. 'gcc': {
  69. 'gen': ['-fprofile-generate', '-fprofile-dir={TEMPDIR}'],
  70. 'use': ['-fprofile-use', '-fprofile-correction', '-fprofile-dir={TEMPDIR}'],
  71. },
  72. # blind copy from 'configure' script in CPython 3.7
  73. 'icc': {
  74. 'gen': ['-prof-gen'],
  75. 'use': ['-prof-use'],
  76. }
  77. }
  78. PGO_CONFIG['mingw32'] = PGO_CONFIG['gcc']
  79. if IS_PY2:
  80. def encode_fs(name):
  81. return name if isinstance(name, bytes) else name.encode(IO_ENCODING)
  82. else:
  83. def encode_fs(name):
  84. return name
  85. @magics_class
  86. class CythonMagics(Magics):
  87. def __init__(self, shell):
  88. super(CythonMagics, self).__init__(shell)
  89. self._reloads = {}
  90. self._code_cache = {}
  91. self._pyximport_installed = False
  92. def _import_all(self, module):
  93. mdict = module.__dict__
  94. if '__all__' in mdict:
  95. keys = mdict['__all__']
  96. else:
  97. keys = [k for k in mdict if not k.startswith('_')]
  98. for k in keys:
  99. try:
  100. self.shell.push({k: mdict[k]})
  101. except KeyError:
  102. msg = "'module' object has no attribute '%s'" % k
  103. raise AttributeError(msg)
  104. @cell_magic
  105. def cython_inline(self, line, cell):
  106. """Compile and run a Cython code cell using Cython.inline.
  107. This magic simply passes the body of the cell to Cython.inline
  108. and returns the result. If the variables `a` and `b` are defined
  109. in the user's namespace, here is a simple example that returns
  110. their sum::
  111. %%cython_inline
  112. return a+b
  113. For most purposes, we recommend the usage of the `%%cython` magic.
  114. """
  115. locs = self.shell.user_global_ns
  116. globs = self.shell.user_ns
  117. return cython_inline(cell, locals=locs, globals=globs)
  118. @cell_magic
  119. def cython_pyximport(self, line, cell):
  120. """Compile and import a Cython code cell using pyximport.
  121. The contents of the cell are written to a `.pyx` file in the current
  122. working directory, which is then imported using `pyximport`. This
  123. magic requires a module name to be passed::
  124. %%cython_pyximport modulename
  125. def f(x):
  126. return 2.0*x
  127. The compiled module is then imported and all of its symbols are
  128. injected into the user's namespace. For most purposes, we recommend
  129. the usage of the `%%cython` magic.
  130. """
  131. module_name = line.strip()
  132. if not module_name:
  133. raise ValueError('module name must be given')
  134. fname = module_name + '.pyx'
  135. with io.open(fname, 'w', encoding='utf-8') as f:
  136. f.write(cell)
  137. if 'pyximport' not in sys.modules or not self._pyximport_installed:
  138. import pyximport
  139. pyximport.install()
  140. self._pyximport_installed = True
  141. if module_name in self._reloads:
  142. module = self._reloads[module_name]
  143. # Note: reloading extension modules is not actually supported
  144. # (requires PEP-489 reinitialisation support).
  145. # Don't know why this should ever have worked as it reads here.
  146. # All we really need to do is to update the globals below.
  147. #reload(module)
  148. else:
  149. __import__(module_name)
  150. module = sys.modules[module_name]
  151. self._reloads[module_name] = module
  152. self._import_all(module)
  153. @magic_arguments.magic_arguments()
  154. @magic_arguments.argument(
  155. '-a', '--annotate', action='store_true', default=False,
  156. help="Produce a colorized HTML version of the source."
  157. )
  158. @magic_arguments.argument(
  159. '-+', '--cplus', action='store_true', default=False,
  160. help="Output a C++ rather than C file."
  161. )
  162. @magic_arguments.argument(
  163. '-3', dest='language_level', action='store_const', const=3, default=None,
  164. help="Select Python 3 syntax."
  165. )
  166. @magic_arguments.argument(
  167. '-2', dest='language_level', action='store_const', const=2, default=None,
  168. help="Select Python 2 syntax."
  169. )
  170. @magic_arguments.argument(
  171. '-f', '--force', action='store_true', default=False,
  172. help="Force the compilation of a new module, even if the source has been "
  173. "previously compiled."
  174. )
  175. @magic_arguments.argument(
  176. '-c', '--compile-args', action='append', default=[],
  177. help="Extra flags to pass to compiler via the `extra_compile_args` "
  178. "Extension flag (can be specified multiple times)."
  179. )
  180. @magic_arguments.argument(
  181. '--link-args', action='append', default=[],
  182. help="Extra flags to pass to linker via the `extra_link_args` "
  183. "Extension flag (can be specified multiple times)."
  184. )
  185. @magic_arguments.argument(
  186. '-l', '--lib', action='append', default=[],
  187. help="Add a library to link the extension against (can be specified "
  188. "multiple times)."
  189. )
  190. @magic_arguments.argument(
  191. '-n', '--name',
  192. help="Specify a name for the Cython module."
  193. )
  194. @magic_arguments.argument(
  195. '-L', dest='library_dirs', metavar='dir', action='append', default=[],
  196. help="Add a path to the list of library directories (can be specified "
  197. "multiple times)."
  198. )
  199. @magic_arguments.argument(
  200. '-I', '--include', action='append', default=[],
  201. help="Add a path to the list of include directories (can be specified "
  202. "multiple times)."
  203. )
  204. @magic_arguments.argument(
  205. '-S', '--src', action='append', default=[],
  206. help="Add a path to the list of src files (can be specified "
  207. "multiple times)."
  208. )
  209. @magic_arguments.argument(
  210. '--pgo', dest='pgo', action='store_true', default=False,
  211. help=("Enable profile guided optimisation in the C compiler. "
  212. "Compiles the cell twice and executes it in between to generate a runtime profile.")
  213. )
  214. @magic_arguments.argument(
  215. '--verbose', dest='quiet', action='store_false', default=True,
  216. help=("Print debug information like generated .c/.cpp file location "
  217. "and exact gcc/g++ command invoked.")
  218. )
  219. @cell_magic
  220. def cython(self, line, cell):
  221. """Compile and import everything from a Cython code cell.
  222. The contents of the cell are written to a `.pyx` file in the
  223. directory `IPYTHONDIR/cython` using a filename with the hash of the
  224. code. This file is then cythonized and compiled. The resulting module
  225. is imported and all of its symbols are injected into the user's
  226. namespace. The usage is similar to that of `%%cython_pyximport` but
  227. you don't have to pass a module name::
  228. %%cython
  229. def f(x):
  230. return 2.0*x
  231. To compile OpenMP codes, pass the required `--compile-args`
  232. and `--link-args`. For example with gcc::
  233. %%cython --compile-args=-fopenmp --link-args=-fopenmp
  234. ...
  235. To enable profile guided optimisation, pass the ``--pgo`` option.
  236. Note that the cell itself needs to take care of establishing a suitable
  237. profile when executed. This can be done by implementing the functions to
  238. optimise, and then calling them directly in the same cell on some realistic
  239. training data like this::
  240. %%cython --pgo
  241. def critical_function(data):
  242. for item in data:
  243. ...
  244. # execute function several times to build profile
  245. from somewhere import some_typical_data
  246. for _ in range(100):
  247. critical_function(some_typical_data)
  248. In Python 3.5 and later, you can distinguish between the profile and
  249. non-profile runs as follows::
  250. if "_pgo_" in __name__:
  251. ... # execute critical code here
  252. """
  253. args = magic_arguments.parse_argstring(self.cython, line)
  254. code = cell if cell.endswith('\n') else cell + '\n'
  255. lib_dir = os.path.join(get_ipython_cache_dir(), 'cython')
  256. key = (code, line, sys.version_info, sys.executable, cython_version)
  257. if not os.path.exists(lib_dir):
  258. os.makedirs(lib_dir)
  259. if args.pgo:
  260. key += ('pgo',)
  261. if args.force:
  262. # Force a new module name by adding the current time to the
  263. # key which is hashed to determine the module name.
  264. key += (time.time(),)
  265. if args.name:
  266. module_name = str(args.name) # no-op in Py3
  267. else:
  268. module_name = "_cython_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest()
  269. html_file = os.path.join(lib_dir, module_name + '.html')
  270. module_path = os.path.join(lib_dir, module_name + self.so_ext)
  271. have_module = os.path.isfile(module_path)
  272. need_cythonize = args.pgo or not have_module
  273. if args.annotate:
  274. if not os.path.isfile(html_file):
  275. need_cythonize = True
  276. extension = None
  277. if need_cythonize:
  278. extensions = self._cythonize(module_name, code, lib_dir, args, quiet=args.quiet)
  279. if extensions is None:
  280. # Compilation failed and printed error message
  281. return None
  282. assert len(extensions) == 1
  283. extension = extensions[0]
  284. self._code_cache[key] = module_name
  285. if args.pgo:
  286. self._profile_pgo_wrapper(extension, lib_dir)
  287. try:
  288. self._build_extension(extension, lib_dir, pgo_step_name='use' if args.pgo else None,
  289. quiet=args.quiet)
  290. except distutils.errors.CompileError:
  291. # Build failed and printed error message
  292. return None
  293. module = imp.load_dynamic(module_name, module_path)
  294. self._import_all(module)
  295. if args.annotate:
  296. try:
  297. with io.open(html_file, encoding='utf-8') as f:
  298. annotated_html = f.read()
  299. except IOError as e:
  300. # File could not be opened. Most likely the user has a version
  301. # of Cython before 0.15.1 (when `cythonize` learned the
  302. # `force` keyword argument) and has already compiled this
  303. # exact source without annotation.
  304. print('Cython completed successfully but the annotated '
  305. 'source could not be read.', file=sys.stderr)
  306. print(e, file=sys.stderr)
  307. else:
  308. return display.HTML(self.clean_annotated_html(annotated_html))
  309. def _profile_pgo_wrapper(self, extension, lib_dir):
  310. """
  311. Generate a .c file for a separate extension module that calls the
  312. module init function of the original module. This makes sure that the
  313. PGO profiler sees the correct .o file of the final module, but it still
  314. allows us to import the module under a different name for profiling,
  315. before recompiling it into the PGO optimised module. Overwriting and
  316. reimporting the same shared library is not portable.
  317. """
  318. extension = copy.copy(extension) # shallow copy, do not modify sources in place!
  319. module_name = extension.name
  320. pgo_module_name = '_pgo_' + module_name
  321. pgo_wrapper_c_file = os.path.join(lib_dir, pgo_module_name + '.c')
  322. with io.open(pgo_wrapper_c_file, 'w', encoding='utf-8') as f:
  323. f.write(textwrap.dedent(u"""
  324. #include "Python.h"
  325. #if PY_MAJOR_VERSION < 3
  326. extern PyMODINIT_FUNC init%(module_name)s(void);
  327. PyMODINIT_FUNC init%(pgo_module_name)s(void); /*proto*/
  328. PyMODINIT_FUNC init%(pgo_module_name)s(void) {
  329. PyObject *sys_modules;
  330. init%(module_name)s(); if (PyErr_Occurred()) return;
  331. sys_modules = PyImport_GetModuleDict(); /* borrowed, no exception, "never" fails */
  332. if (sys_modules) {
  333. PyObject *module = PyDict_GetItemString(sys_modules, "%(module_name)s"); if (!module) return;
  334. PyDict_SetItemString(sys_modules, "%(pgo_module_name)s", module);
  335. Py_DECREF(module);
  336. }
  337. }
  338. #else
  339. extern PyMODINIT_FUNC PyInit_%(module_name)s(void);
  340. PyMODINIT_FUNC PyInit_%(pgo_module_name)s(void); /*proto*/
  341. PyMODINIT_FUNC PyInit_%(pgo_module_name)s(void) {
  342. return PyInit_%(module_name)s();
  343. }
  344. #endif
  345. """ % {'module_name': module_name, 'pgo_module_name': pgo_module_name}))
  346. extension.sources = extension.sources + [pgo_wrapper_c_file] # do not modify in place!
  347. extension.name = pgo_module_name
  348. self._build_extension(extension, lib_dir, pgo_step_name='gen')
  349. # import and execute module code to generate profile
  350. so_module_path = os.path.join(lib_dir, pgo_module_name + self.so_ext)
  351. imp.load_dynamic(pgo_module_name, so_module_path)
  352. def _cythonize(self, module_name, code, lib_dir, args, quiet=True):
  353. pyx_file = os.path.join(lib_dir, module_name + '.pyx')
  354. pyx_file = encode_fs(pyx_file)
  355. c_include_dirs = args.include
  356. c_src_files = list(map(str, args.src))
  357. if 'numpy' in code:
  358. import numpy
  359. c_include_dirs.append(numpy.get_include())
  360. with io.open(pyx_file, 'w', encoding='utf-8') as f:
  361. f.write(code)
  362. extension = Extension(
  363. name=module_name,
  364. sources=[pyx_file] + c_src_files,
  365. include_dirs=c_include_dirs,
  366. library_dirs=args.library_dirs,
  367. extra_compile_args=args.compile_args,
  368. extra_link_args=args.link_args,
  369. libraries=args.lib,
  370. language='c++' if args.cplus else 'c',
  371. )
  372. try:
  373. opts = dict(
  374. quiet=quiet,
  375. annotate=args.annotate,
  376. force=True,
  377. )
  378. if args.language_level is not None:
  379. assert args.language_level in (2, 3)
  380. opts['language_level'] = args.language_level
  381. elif sys.version_info[0] >= 3:
  382. opts['language_level'] = 3
  383. return cythonize([extension], **opts)
  384. except CompileError:
  385. return None
  386. def _build_extension(self, extension, lib_dir, temp_dir=None, pgo_step_name=None, quiet=True):
  387. build_extension = self._get_build_extension(
  388. extension, lib_dir=lib_dir, temp_dir=temp_dir, pgo_step_name=pgo_step_name)
  389. old_threshold = None
  390. try:
  391. if not quiet:
  392. old_threshold = distutils.log.set_threshold(distutils.log.DEBUG)
  393. build_extension.run()
  394. finally:
  395. if not quiet and old_threshold is not None:
  396. distutils.log.set_threshold(old_threshold)
  397. def _add_pgo_flags(self, build_extension, step_name, temp_dir):
  398. compiler_type = build_extension.compiler.compiler_type
  399. if compiler_type == 'unix':
  400. compiler_cmd = build_extension.compiler.compiler_so
  401. # TODO: we could try to call "[cmd] --version" for better insights
  402. if not compiler_cmd:
  403. pass
  404. elif 'clang' in compiler_cmd or 'clang' in compiler_cmd[0]:
  405. compiler_type = 'clang'
  406. elif 'icc' in compiler_cmd or 'icc' in compiler_cmd[0]:
  407. compiler_type = 'icc'
  408. elif 'gcc' in compiler_cmd or 'gcc' in compiler_cmd[0]:
  409. compiler_type = 'gcc'
  410. elif 'g++' in compiler_cmd or 'g++' in compiler_cmd[0]:
  411. compiler_type = 'gcc'
  412. config = PGO_CONFIG.get(compiler_type)
  413. orig_flags = []
  414. if config and step_name in config:
  415. flags = [f.format(TEMPDIR=temp_dir) for f in config[step_name]]
  416. for extension in build_extension.extensions:
  417. orig_flags.append((extension.extra_compile_args, extension.extra_link_args))
  418. extension.extra_compile_args = extension.extra_compile_args + flags
  419. extension.extra_link_args = extension.extra_link_args + flags
  420. else:
  421. print("No PGO %s configuration known for C compiler type '%s'" % (step_name, compiler_type),
  422. file=sys.stderr)
  423. return orig_flags
  424. @property
  425. def so_ext(self):
  426. """The extension suffix for compiled modules."""
  427. try:
  428. return self._so_ext
  429. except AttributeError:
  430. self._so_ext = self._get_build_extension().get_ext_filename('')
  431. return self._so_ext
  432. def _clear_distutils_mkpath_cache(self):
  433. """clear distutils mkpath cache
  434. prevents distutils from skipping re-creation of dirs that have been removed
  435. """
  436. try:
  437. from distutils.dir_util import _path_created
  438. except ImportError:
  439. pass
  440. else:
  441. _path_created.clear()
  442. def _get_build_extension(self, extension=None, lib_dir=None, temp_dir=None,
  443. pgo_step_name=None, _build_ext=build_ext):
  444. self._clear_distutils_mkpath_cache()
  445. dist = Distribution()
  446. config_files = dist.find_config_files()
  447. try:
  448. config_files.remove('setup.cfg')
  449. except ValueError:
  450. pass
  451. dist.parse_config_files(config_files)
  452. if not temp_dir:
  453. temp_dir = lib_dir
  454. add_pgo_flags = self._add_pgo_flags
  455. if pgo_step_name:
  456. base_build_ext = _build_ext
  457. class _build_ext(_build_ext):
  458. def build_extensions(self):
  459. add_pgo_flags(self, pgo_step_name, temp_dir)
  460. base_build_ext.build_extensions(self)
  461. build_extension = _build_ext(dist)
  462. build_extension.finalize_options()
  463. if temp_dir:
  464. temp_dir = encode_fs(temp_dir)
  465. build_extension.build_temp = temp_dir
  466. if lib_dir:
  467. lib_dir = encode_fs(lib_dir)
  468. build_extension.build_lib = lib_dir
  469. if extension is not None:
  470. build_extension.extensions = [extension]
  471. return build_extension
  472. @staticmethod
  473. def clean_annotated_html(html):
  474. """Clean up the annotated HTML source.
  475. Strips the link to the generated C or C++ file, which we do not
  476. present to the user.
  477. """
  478. r = re.compile('<p>Raw output: <a href="(.*)">(.*)</a>')
  479. html = '\n'.join(l for l in html.splitlines() if not r.match(l))
  480. return html
  481. __doc__ = __doc__.format(
  482. # rST doesn't see the -+ flag as part of an option list, so we
  483. # hide it from the module-level docstring.
  484. CYTHON_DOC=dedent(CythonMagics.cython.__doc__\
  485. .replace('-+, --cplus', '--cplus ')),
  486. CYTHON_INLINE_DOC=dedent(CythonMagics.cython_inline.__doc__),
  487. CYTHON_PYXIMPORT_DOC=dedent(CythonMagics.cython_pyximport.__doc__),
  488. )