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