Subclassing Profile.Profile to use it as a context manager

Question

I was profiling my code and annoyed by the continuous calls of

Profile.run(my_function())
Profile.run(my_next_function())

so I decided to make a context manager instead. This seems pretty straightforward but I wanted to make sure nothing about it was bad - I was a little surprised that you couldn't use it as a context manager to begin with which makes me suspect someone had a reason to not implement it.

try:
    import cProfile as Profile
except ImportError:
    import profile as Profile
finally:
    import pstats
    import sys


class context_profiler(Profile.Profile):

    def __init__(self, sortby='cumulative', stream=sys.stdout, **kwargs):
        super(context_profiler, self).__init__(**kwargs)
        self.sortby = sortby
        self.stream = stream

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

    def __exit__(self, exc_type, exc_value, traceback):
        self.disable()
        self.print_data()

    def print_data(self):
        ps = pstats.Stats(self, stream=self.stream).sort_stats(self.sortby)
        ps.print_stats()


if __name__ == '__main__':
    with context_profiler() as pr:
        for i in xrange(2):
            for j in range(10000):
                pr.stream.write(str(i+j))

Answer 1

An __exit__ method should analyze the exc_type parameter. If there really was an exception, the collected profile is most likely not representative, and should not be accumulated. At least provide an option for such behaviour.

Otherwise looks pretty good.

Answer 2

Remember all that you've learned in the last year, in particular about naming. You probably want to name your class using PascalCasing.

Also, instead of breaking encapsulation and saying pr.stream.write you should have your class implement write

def write(self, string):
    self.stream.write(string)

You may also want to implement flush

def flush(self):
    self.stream.flush()