7
\$\begingroup\$

I have created a simple profiler in Python, but I feel like there are better ways to do some things.

Terminology:

  • Section - a block of code to profile
  • Label - A name for a section to be used as a key in profiler.info
  • Label-specific time: The time spent in a label excluding time spent in labels inside it.
  • Total time: The time spent in a label including time spent in labels inside it.

The code:

from collections import defaultdict
from time import time, sleep
from contextlib import contextmanager
from functools import wraps

class SecInfo():
  """
  Container for section info.

  Attributes:
    entered:
      The number of times the section has been entered.

    time:
      The amount of time spent in the section, not counting sections inside it.

    total time:
      The amount of time spent in the section, counting sections inside it.
  """
  def __init__(self):
    """
    Create a SecInfo object. Initializes everthing to 0.
    """
    self.entered = 0
    self.time = 0
    self.total_time = 0

class Profiler():
  """
  A profiler.

  Attributes:
    info: A defaultdict of SecInfos for every section that has been profiled.


  Methods:
    reset()
    results()
    start_section(label)
    end_section()
    profile(func)
    profile_label(label)
    ctx(label)
  """
  def __init__(self):
    """
    Create a Profiler. Initializes private variables and info.
    """
    self._curr_time = None # initialized when profiling starts
    self._all_starts = {}
    self._curr_labels = []
    self.info = defaultdict(SecInfo)

  def reset(self):
    """
    Reset info to empty defaultdict.
    """
    self.info.clear()

  def results(self):
    """
    Print human-readable version of info.
    """
    for label, info in self.info.items():
      print(f"{label}:")
      print(f"Entered {info.entered} time(s)")
      print(f"Label-specific time is {info.time:.15f}")
      print(f"Total time is {info.total_time:.15f}")
      print()

  def start_section(self, label):
    """
    Start a section with label as the label.

    Arguments:
      label:
        The label to use for the section.
    """
    t = time()
    if label not in self._all_starts: # keep track of only first recursive start time
      self._all_starts[label] = t

    if self._curr_labels:
      self.info[self._curr_labels[-1]].time += t - self._curr_time

    self._curr_time = t
    self._curr_labels.append(label)
    self.info[label].entered += 1

  def end_section(self):
    """
    End the current section. If sections are nested, end the most recent.
    """
    t = time()
    label = self._curr_labels.pop()
    info = self.info[label]
    info.time += t - self._curr_time

    if label not in self._curr_labels: # if it's not a recursive call
      info.total_time += t - self._all_starts.pop(label)

    self._curr_time = t

  def _profile(self, func, label):
    """
    Generates a wrapper to profile func with the given label.

    Arguments:
      func:
        The function to wrap

      label:
        The label to use
    """
    @wraps(func)
    def wrapper(*args):
      self.start_section(label)
      result = func(*args)
      self.end_section()
      return result

    return wrapper

  def profile(self, func):
    """
    A decorator. Profiles the decorated function, with label func.__name__.
    """

    return self._profile(func, func.__name__)

  def profile_label(self, label):
    """
    A decorator generator. Takes a label, and generates a decorator to profile
    the wrapped function with the given label.

    Arguments:
      label:
        The label to profile with.
    """
    def deco(func):
      """
      A decorator. Profiles the given function with the label passed to
      profile.profile_label.
      """
      return self._profile(func, label)

    return deco

  @contextmanager
  def ctx(self, label):
    """
    Use like
      with profiler.ctx("label"):
        # stuff to profile

    . Identical to wrapping block with
    start_section("label") and end_section().
    """
    self.start_section(label)
    yield self
    self.end_section()

Example usage/testing (comments are label: specific, total):

p = Profiler()

@p.profile
def g():
  sleep(2)

p.start_section("g_call")
g()
p.end_section()

p.results() # g_call: 0, 2; g: 2, 2
p.reset()

print()

@p.profile
def f():
  g()

with p.ctx("f_call"):
  f()

p.results() # f_call: 0, 2; f: 0, 2; g: 2, 2
p.reset()

print()

@p.profile
def recursive(n):
  if n != 0:
    sleep(1)
    recursive(n - 1)

recursive(5)

p.results() # recursive: 5, 5
p.reset()

print()

@p.profile
def recursive1(n):
  if n != 0:
    sleep(1)
    recursive2(n - 1)

@p.profile
def recursive2(n):
  if n != 0:
    sleep(1)
    recursive1(n - 1)

recursive1(5)

p.results() # recursive1: 3, 5; recursive2: 2, 4

Specific concerns:

  • In profile_label, should I be using a lambda instead?
  • Is there a simpler way to keep track of label-specific time?
  • Are there any bugs, especially recursion related?
\$\endgroup\$
2
\$\begingroup\$

I would make more use of your context manager. While it is clear from the docstring of end_section that it ends the most recent section, this is not so clear just from the code. Here the Contextmanager has the big advantage that it creates an indented scope, so it makes this obvious.

I would also rename ctx to something more readable, maybe just section?

You can use this method in your _profile method.

Finally, I would make start_section and end_section protected by adding a _ at the beginning of the name, signalling to the user that it is discouraged to use it (and that they should use the Contextmanager instead).

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.