controller-wisetopdataserver/python/biopro/util/profile.py

import abc
import io
import os
from subprocess import Popen

from .cli_base import CliOptions, cli_options
from .stack import print_exception
from .sys import PID, current_command


# noinspection PyUnusedLocal
class ProfileOptions(CliOptions):
    def __init__(self):
        self.profiling_method = None
        self.profiling_output = None
        self.profiling_sorting = ('cumtime',)
        self.profiling_output_type = 'text'

    @cli_options('--profile', value='METHOD', optional=True)
    def _profile_method(self, opt: str, value: str):
        """enable profiling and choose profile method. could be:
        profile  : use cProfile
        yappi    : use yappi (default, require yappi)
        """
        if value is None:
            self.profiling_method = 'default'
        else:

            if value not in ('cprofile', 'yappi'):
                raise RuntimeError('illegal profile method : ' + value)

            self.profiling_method = value

    @cli_options('--profile-output', value='PATH')
    def _profile_output(self, opt: str, value: str):
        """profiling file output path"""
        self.profiling_output = value

    @cli_options('--profile-output-type', value='FORMAT')
    def _profile_output_type(self, opt: str, value: str):
        """profiling file format. could be:
        raw  :
        text : (default)
        dot  : (required gprof2dot)"""

        if value not in ('raw', 'text', 'dot'):
            raise RuntimeError('illegal output type : ' + value)

        self.profiling_output_type = value

    @cli_options('--profile-sorting', value='FIELDs')
    def _profile_sorting(self, opt: str, value: str):
        """the sorting method of the field in profiling state.
        detail : https://docs.python.org/3.5/library/profile.html#pstats.Stats.sort_stats"""

        fields = tuple(value.split(','))

        for field in fields:

            if field not in ('calls', 'cumulative', 'cumtime', 'file', 'filename', 'module',
                             'ncalls', 'pcalls', 'line', 'name', 'nfl', 'stdname', 'time', 'tottime'):
                raise RuntimeError('illegal field : ' + field)

        self.profiling_sorting = fields

    def generate_profile_context(self) -> 'ProfileContext':
        if self.profiling_method is None:
            ret = _NullProfileContext(self)

        else:
            try:
                if self.profiling_method == 'cprofile':
                    ret = _CProfileContext(self)
                else:
                    try:
                        if self.profiling_method == 'yappi':
                            ret = _YappiProfileContext(self)
                        else:
                            ret = _YappiProfileContext(self)
                    except ImportError:
                        print_exception()
                        ret = _CProfileContext(self)
            except ImportError:
                print_exception()
                ret = _NullProfileContext(self)

        return ret


class ProfileContext(metaclass=abc.ABCMeta):

    def __init__(self, options: ProfileOptions):
        self._method = options.profiling_method
        self._output = options.profiling_output
        self._sorting = options.profiling_sorting
        self._output_type = options.profiling_output_type

    def __enter__(self):
        self.setup()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            print_exception(exc_info=(exc_type, exc_val, exc_tb))

        try:
            self.finalize()
        except:
            print_exception()

        return True

    @abc.abstractmethod
    def setup(self):
        pass

    @abc.abstractmethod
    def finalize(self):
        pass

    def dump_pstat(self, ps):
        if self._output_type == 'raw' or self._output_type == 'dot':
            profiling_output = self._output

            if profiling_output is None:
                profiling_output = os.environ.get("HOME") + "/.bps-%d.profile" % PID

            ps.dump_stats(profiling_output)

            print('dump', profiling_output)

            if self._output_type == 'dot':
                self_cmd = current_command()[0]

                output_dot_path = profiling_output + '.png'
                cmd_line = '{} -m gprof2dot -f pstats {} | dot -T png -o {}'.format(self_cmd,
                                                                                    profiling_output,
                                                                                    output_dot_path)

                cmd = ['bash', '-c', cmd_line]

                with Popen(cmd) as process:
                    process.wait()

                print('dump', output_dot_path)

        elif self._output_type == 'text':
            if self._output is not None:
                out = open(self._output, 'w')
            else:
                out = io.StringIO()

            try:
                ps.stream = out
                ps.sort_stats(*self._sorting)
                ps.print_stats()

            finally:
                if self._output is not None:
                    out.close()
                else:
                    print(out.getvalue())


class _NullProfileContext(ProfileContext):

    def __init__(self, options: ProfileOptions):
        super().__init__(options)

    def setup(self):
        pass

    def finalize(self):
        pass


class _CProfileContext(ProfileContext):

    def __init__(self, options: ProfileOptions):
        super().__init__(options)

        import cProfile

        self._profile = cProfile.Profile()

    def setup(self):
        self._profile.enable()

    def finalize(self):
        self._profile.disable()

        import pstats

        self.dump_pstat(pstats.Stats(self._profile))


class _YappiProfileContext(ProfileContext):
    def __init__(self, options: ProfileOptions):
        super().__init__(options)

        import yappi

        self._yappi = yappi

    def setup(self):
        self._yappi.start()

    def finalize(self):
        self._yappi.stop()

        stats = self._yappi.get_func_stats()
        ps = self._yappi.convert2pstats(stats)

        self.dump_pstat(ps)