#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import absolute_import

__version__ = '3.0.2'

try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO

from IPython.core.magic import (Magics, magics_class, line_magic)
from IPython.core.page import page
from IPython.utils.ipstruct import Struct
from IPython.core.error import UsageError

@magics_class
class LineProfilerMagics(Magics):

    @line_magic
    def lprun(self, parameter_s=''):
        """ Execute a statement under the line-by-line profiler from the
        line_profiler module.

        Usage:
          %lprun -f func1 -f func2 <statement>

        The given statement (which doesn't require quote marks) is run via the
        LineProfiler. Profiling is enabled for the functions specified by the -f
        options. The statistics will be shown side-by-side with the code through the
        pager once the statement has completed.

        Options:

        -f <function>: LineProfiler only profiles functions and methods it is told
        to profile.  This option tells the profiler about these functions. Multiple
        -f options may be used. The argument may be any expression that gives
        a Python function or method object. However, one must be careful to avoid
        spaces that may confuse the option parser.

        -m <module>: Get all the functions/methods in a module

        One or more -f or -m options are required to get any useful results.

        -D <filename>: dump the raw statistics out to a pickle file on disk. The
        usual extension for this is ".lprof". These statistics may be viewed later
        by running line_profiler.py as a script.

        -T <filename>: dump the text-formatted statistics with the code side-by-side
        out to a text file.

        -r: return the LineProfiler object after it has completed profiling.

        -s: strip out all entries from the print-out that have zeros.

        -u: specify metric unit for the print-out in seconds.
        """

        # Escape quote markers.
        opts_def = Struct(D=[''], T=[''], f=[], m=[], u=None)
        parameter_s = parameter_s.replace('"', r'\"').replace("'", r"\'")
        opts, arg_str = self.parse_options(parameter_s, 'rsf:m:D:T:u:', list_all=True)
        opts.merge(opts_def)

        global_ns = self.shell.user_global_ns
        local_ns = self.shell.user_ns

        # Get the requested functions.
        funcs = []
        for name in opts.f:
            try:
                funcs.append(eval(name, global_ns, local_ns))
            except Exception as e:
                raise UsageError('Could not find function %r.\n%s: %s' % (name,
                    e.__class__.__name__, e))

        profile = LineProfiler(*funcs)

        # Get the modules, too
        for modname in opts.m:
            try:
                mod = __import__(modname, fromlist=[''])
                profile.add_module(mod)
            except Exception as e:
                raise UsageError('Could not find module %r.\n%s: %s' % (modname,
                    e.__class__.__name__, e))

        if opts.u is not None:
            try:
                output_unit = float(opts.u[0])
            except Exception as e:
                raise TypeError("Metric unit setting must be a float.")
        else:
            output_unit = None

        # Add the profiler to the builtins for @profile.
        if PY3:
            import builtins
        else:
            import __builtin__ as builtins

        if 'profile' in builtins.__dict__:
            had_profile = True
            old_profile = builtins.__dict__['profile']
        else:
            had_profile = False
            old_profile = None
        builtins.__dict__['profile'] = profile

        try:
            try:
                profile.runctx(arg_str, global_ns, local_ns)
                message = ''
            except SystemExit:
                message = """*** SystemExit exception caught in code being profiled."""
            except KeyboardInterrupt:
                message = ("*** KeyboardInterrupt exception caught in code being "
                    "profiled.")
        finally:
            if had_profile:
                builtins.__dict__['profile'] = old_profile

        # Trap text output.
        stdout_trap = StringIO()
        profile.print_stats(stdout_trap, output_unit=output_unit, stripzeros='s' in opts)
        output = stdout_trap.getvalue()
        output = output.rstrip()

        page(output)
        print(message, end="")

        dump_file = opts.D[0]
        if dump_file:
            profile.dump_stats(dump_file)
            print('\n*** Profile stats pickled to file %r. %s' % (
                dump_file, message))

        text_file = opts.T[0]
        if text_file:
            pfile = open(text_file, 'w')
            pfile.write(output)
            pfile.close()
            print('\n*** Profile printout saved to text file %r. %s' % (
                text_file, message))

        return_value = None
        if 'r' in opts:
            return_value = profile

        return return_value

