PR: Add profile, profilefile, profilecell magics

spyder_kernels/customize/code_runner.py

@@ -14,10 +14,13 @@
 import bdb
 import builtins
 from contextlib import contextmanager
+import cProfile
+from functools import partial
 import io
 import logging
 import os
 import pdb
+import tempfile
 import shlex
 import sys
 import time
@@ -30,18 +33,21 @@
 )
 from IPython.core.magic import (
     needs_local_scope,
+    no_var_expand,
+    line_cell_magic,
     magics_class,
     Magics,
     line_magic,
 )
 from IPython.core import magic_arguments
 
 # Local imports
-from spyder_kernels.comms.frontendcomm import frontend_request
+from spyder_kernels.comms.frontendcomm import frontend_request, CommError
 from spyder_kernels.customize.namespace_manager import NamespaceManager
 from spyder_kernels.customize.spyderpdb import SpyderPdb
 from spyder_kernels.customize.umr import UserModuleReloader
-from spyder_kernels.customize.utils import capture_last_Expr, canonic
+from spyder_kernels.customize.utils import (
+    capture_last_Expr, canonic, create_pathlist)
 
 
 # For logging
@@ -192,6 +198,27 @@ def debugfile(self, line, local_ns=None):
                 context_locals=local_ns,
             )
 
+    @runfile_arguments
+    @needs_local_scope
+    @line_magic
+    def profilefile(self, line, local_ns=None):
+        """Profile a file."""
+        args, local_ns = self._parse_runfile_argstring(
+            self.profilefile, line, local_ns)
+
+        with self._profile_exec() as prof_exec:
+            self._exec_file(
+                filename=args.filename,
+                canonic_filename=args.canonic_filename,
+                wdir=args.wdir,
+                current_namespace=args.current_namespace,
+                args=args.args,
+                exec_fun=prof_exec,
+                post_mortem=args.post_mortem,
+                context_globals=args.namespace,
+                context_locals=local_ns,
+            )
+
     @runcell_arguments
     @needs_local_scope
     @line_magic
@@ -230,6 +257,34 @@ def debugcell(self, line, local_ns=None):
                 context_locals=local_ns,
             )
 
+    @runcell_arguments
+    @needs_local_scope
+    @line_magic
+    def profilecell(self, line, local_ns=None):
+        """Profile a code cell."""
+        args = self._parse_runcell_argstring(self.profilecell, line)
+
+        with self._profile_exec() as prof_exec:
+            return self._exec_cell(
+                cell_id=args.cell_id,
+                filename=args.filename,
+                canonic_filename=args.canonic_filename,
+                exec_fun=prof_exec,
+                post_mortem=args.post_mortem,
+                context_globals=self.shell.user_ns,
+                context_locals=local_ns,
+            )
+
+    @no_var_expand
+    @needs_local_scope
+    @line_cell_magic
+    def profile(self, line, cell=None, local_ns=None):
+        """Profile the given line."""
+        if cell is not None:
+            line += "\n" + cell
+        with self._profile_exec() as prof_exec:
+            return prof_exec(line, self.shell.user_ns, local_ns)
+
     @contextmanager
     def _debugger_exec(self, filename, continue_if_has_breakpoints):
         """Get an exec function to use for debugging."""
@@ -251,6 +306,57 @@ def debug_exec(code, glob, loc):
             # Enter recursive debugger
             yield debug_exec
 
+    @contextmanager
+    def _profile_exec(self):
+        """Get an exec function for profiling."""
+        tmp_dir = None
+        if sys.platform.startswith('linux'):
+            # Do not use /tmp for temporary files
+            try:
+                from xdg.BaseDirectory import xdg_data_home
+                tmp_dir = xdg_data_home
+                os.makedirs(tmp_dir, exist_ok=True)
+            except Exception:
+                tmp_dir = None
+        with tempfile.TemporaryDirectory(dir=tmp_dir) as tempdir:
+            # Reset the tracing function in case we are debugging
+            trace_fun = sys.gettrace()
+            sys.settrace(None)
+
+            # Get a file to save the results
+            profile_filename = os.path.join(tempdir, "profile.prof")
+
+            try:
+                if self.shell.is_debugging():
+                    def prof_exec(code, glob, loc):
+                        """
+                        If we are debugging (tracing), call_tracing is
+                        necessary for profiling.
+                        """
+                        return sys.call_tracing(cProfile.runctx, (
+                            code, glob, loc, profile_filename
+                        ))
+
+                    yield prof_exec
+                else:
+                    yield partial(cProfile.runctx, filename=profile_filename)
+            finally:
+                # Reset tracing function
+                sys.settrace(trace_fun)
+
+                # Send result to frontend
+                if os.path.isfile(profile_filename):
+                    with open(profile_filename, "br") as f:
+                        profile_result = f.read()
+                    try:
+                        frontend_request(blocking=False).show_profile_file(
+                            profile_result, create_pathlist()
+                        )
+                    except CommError:
+                        logger.debug(
+                            "Could not send profile result to the frontend."
+                        )
+
     def _exec_file(
         self,
         filename=None,