Merge pull request #2225 from bfroehle/cython_annotate

Add "--annotate" option to `%%cython` magic.

IPython/extensions/cythonmagic.py

@@ -15,9 +15,14 b' Parts of this code were taken from Cython.inline.'
 # The full license is in the file COPYING.txt, distributed with this software.
 #-----------------------------------------------------------------------------
 
-import io
-import os, sys
+from __future__ import print_function
+
 import imp
+import io
+import os
+import re
+import sys
+import time
 
 try:
     import hashlib
@@ -27,16 +32,14 b' except ImportError:'
 from distutils.core import Distribution, Extension
 from distutils.command.build_ext import build_ext
 
+from IPython.core import display
+from IPython.core import magic_arguments
 from IPython.core.magic import Magics, magics_class, cell_magic
 from IPython.testing.skipdoctest import skip_doctest
-from IPython.core.magic_arguments import (
-    argument, magic_arguments, parse_argstring
-)
 from IPython.utils import py3compat
 
 import Cython
 from Cython.Compiler.Errors import CompileError
-from Cython.Compiler.Main import Context, default_options
 from Cython.Build.Dependencies import cythonize
 
 
@@ -61,7 +64,7 b' class CythonMagics(Magics):'
         and returns the result. If the variables `a` and `b` are defined
         in the user's namespace, here is a simple example that returns
         their sum::
-        
+
             %%cython_inline
             return a+b
 
@@ -78,14 +81,14 b' class CythonMagics(Magics):'
         The contents of the cell are written to a `.pyx` file in the current
         working directory, which is then imported using `pyximport`. This
         magic requires a module name to be passed::
-        
+
             %%cython_pyximport modulename
             def f(x):
                 return 2.0*x
 
-        The compiled module is then imported and all of its symbols are injected into
-        the user's namespace. For most purposes, we recommend the usage of the
-        `%%cython` magic.
+        The compiled module is then imported and all of its symbols are
+        injected into the user's namespace. For most purposes, we recommend
+        the usage of the `%%cython` magic.
         """
         module_name = line.strip()
         if not module_name:
@@ -105,30 +108,47 b' class CythonMagics(Magics):'
             self._reloads[module_name] = module
         self._import_all(module)
 
-    @magic_arguments()
-    @argument(
+    @magic_arguments.magic_arguments()
+    @magic_arguments.argument(
         '-c', '--compile-args', action='append', default=[],
-        help="Extra flags to pass to compiler via the `extra_compile_args` Extension flag (can be specified  multiple times)."
+        help="Extra flags to pass to compiler via the `extra_compile_args` "
+             "Extension flag (can be specified  multiple times)."
+    )
+    @magic_arguments.argument(
+        '-la', '--link-args', action='append', default=[],
+        help="Extra flags to pass to linker via the `extra_link_args` "
+             "Extension flag (can be specified  multiple times)."
     )
-    @argument(
+    @magic_arguments.argument(
         '-l', '--lib', action='append', default=[],
-        help="Add a library to link the extension against (can be specified  multiple times)."
+        help="Add a library to link the extension against (can be specified "
+             "multiple times)."
     )
-    @argument(
+    @magic_arguments.argument(
         '-I', '--include', action='append', default=[],
-        help="Add a path to the list of include directories (can be specified  multiple times)."
+        help="Add a path to the list of include directories (can be specified "
+             "multiple times)."
     )
-    @argument(
+    @magic_arguments.argument(
+        '-+', '--cplus', action='store_true', default=False,
+        help="Output a C++ rather than C file."
+    )
+    @magic_arguments.argument(
         '-f', '--force', action='store_true', default=False,
-        help="Force the compilation of the pyx module even if it hasn't changed"
+        help="Force the compilation of a new module, even if the source has been "
+             "previously compiled."
+    )
+    @magic_arguments.argument(
+        '-a', '--annotate', action='store_true', default=False,
+        help="Produce a colorized HTML version of the source."
     )
     @cell_magic
     def cython(self, line, cell):
         """Compile and import everything from a Cython code cell.
 
         The contents of the cell are written to a `.pyx` file in the
-        directory `IPYTHONDIR/cython` using a filename with the hash of the code.
-        This file is then cythonized and compiled. The resulting module
+        directory `IPYTHONDIR/cython` using a filename with the hash of the
+        code. This file is then cythonized and compiled. The resulting module
         is imported and all of its symbols are injected into the user's
         namespace. The usage is similar to that of `%%cython_pyximport` but
         you don't have to pass a module name::
@@ -137,22 +157,32 b' class CythonMagics(Magics):'
         def f(x):
             return 2.0*x
         """
-        args = parse_argstring(self.cython, line)
+        args = magic_arguments.parse_argstring(self.cython, line)
         code = cell if cell.endswith('\n') else cell+'\n'
         lib_dir = os.path.join(self.shell.ipython_dir, 'cython')
-        cython_include_dirs = ['.']
-        force = args.force
         quiet = True
-        ctx = Context(cython_include_dirs, default_options)
         key = code, sys.version_info, sys.executable, Cython.__version__
-        module_name = "_cython_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest()
-        so_ext = [ ext for ext,_,mod_type in imp.get_suffixes() if mod_type == imp.C_EXTENSION ][0]
-        module_path = os.path.join(lib_dir, module_name+so_ext)
 
         if not os.path.exists(lib_dir):
             os.makedirs(lib_dir)
 
-        if force or not os.path.isfile(module_path):
+        if args.force:
+            # Force a new module name by adding the current time to the
+            # key which is hashed to determine the module name.
+            key += time.time(),
+
+        module_name = "_cython_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest()
+        module_path = os.path.join(lib_dir, module_name + self.so_ext)
+
+        have_module = os.path.isfile(module_path)
+        need_cythonize = not have_module
+
+        if args.annotate:
+            html_file = os.path.join(lib_dir, module_name + '.html')
+            if not os.path.isfile(html_file):
+                need_cythonize = True
+
+        if need_cythonize:
             c_include_dirs = args.include
             if 'numpy' in code:
                 import numpy
@@ -166,21 +196,22 b' class CythonMagics(Magics):'
                 sources = [pyx_file],
                 include_dirs = c_include_dirs,
                 extra_compile_args = args.compile_args,
+                extra_link_args = args.link_args,
                 libraries = args.lib,
+                language = 'c++' if args.cplus else 'c',
             )
-            dist = Distribution()
-            config_files = dist.find_config_files()
-            try: 
-                config_files.remove('setup.cfg')
-            except ValueError:
-                pass
-            dist.parse_config_files(config_files)
-            build_extension = build_ext(dist)
-            build_extension.finalize_options()
+            build_extension = self._get_build_extension()
             try:
-                build_extension.extensions = cythonize([extension], ctx=ctx, quiet=quiet)
+                opts = dict(
+                    quiet=quiet,
+                    annotate = args.annotate,
+                    force = True,
+                    )
+                build_extension.extensions = cythonize([extension], **opts)
             except CompileError:
                 return
+
+        if not have_module:
             build_extension.build_temp = os.path.dirname(pyx_file)
             build_extension.build_lib  = lib_dir
             build_extension.run()
@@ -189,6 +220,52 b' class CythonMagics(Magics):'
         module = imp.load_dynamic(module_name, module_path)
         self._import_all(module)
 
+        if args.annotate:
+            try:
+                with io.open(html_file, encoding='utf-8') as f:
+                    annotated_html = f.read()
+            except IOError as e:
+                # File could not be opened. Most likely the user has a version
+                # of Cython before 0.15.1 (when `cythonize` learned the
+                # `force` keyword argument) and has already compiled this
+                # exact source without annotation.
+                print('Cython completed successfully but the annotated '
+                      'source could not be read.', file=sys.stderr)
+                print(e, file=sys.stderr)
+            else:
+                return display.HTML(self.clean_annotated_html(annotated_html))
+
+    @property
+    def so_ext(self):
+        """The extension suffix for compiled modules."""
+        try:
+            return self._so_ext
+        except AttributeError:
+            self._so_ext = self._get_build_extension().get_ext_filename('')
+            return self._so_ext
+
+    def _get_build_extension(self):
+        dist = Distribution()
+        config_files = dist.find_config_files()
+        try:
+            config_files.remove('setup.cfg')
+        except ValueError:
+            pass
+        dist.parse_config_files(config_files)
+        build_extension = build_ext(dist)
+        build_extension.finalize_options()
+        return build_extension
+
+    @staticmethod
+    def clean_annotated_html(html):
+        """Clean up the annotated HTML source.
+
+        Strips the link to the generated C or C++ file, which we do not
+        present to the user.
+        """
+        r = re.compile('<p>Raw output: <a href="(.*)">(.*)</a>')
+        html = '\n'.join(l for l in html.splitlines() if not r.match(l))
+        return html
 
 _loaded = False