1# Copyright 2016 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4import contextlib
5import inspect
6import time
7import functools
8
9import log
10from py_trace_event import trace_time
11
12
13@contextlib.contextmanager
14def trace(name, **kwargs):
15  category = "python"
16  start = trace_time.Now()
17  args_to_log = {key: repr(value) for key, value in kwargs.iteritems()}
18  log.add_trace_event("B", start, category, name, args_to_log)
19  try:
20    yield
21  finally:
22    end = trace_time.Now()
23    log.add_trace_event("E", end, category, name)
24
25def traced(*args):
26  def get_wrapper(func):
27    if inspect.isgeneratorfunction(func):
28      raise Exception("Can not trace generators.")
29
30    category = "python"
31
32    arg_spec = inspect.getargspec(func)
33    is_method = arg_spec.args and arg_spec.args[0] == "self"
34
35    def arg_spec_tuple(name):
36      arg_index = arg_spec.args.index(name)
37      defaults_length = len(arg_spec.defaults) if arg_spec.defaults else 0
38      default_index = arg_index + defaults_length - len(arg_spec.args)
39      if default_index >= 0:
40        default = arg_spec.defaults[default_index]
41      else:
42        default = None
43      return (name, arg_index, default)
44
45    args_to_log = map(arg_spec_tuple, arg_names)
46
47    @functools.wraps(func)
48    def traced_function(*args, **kwargs):
49      # Everything outside traced_function is done at decoration-time.
50      # Everything inside traced_function is done at run-time and must be fast.
51      if not log._enabled:  # This check must be at run-time.
52        return func(*args, **kwargs)
53
54      def get_arg_value(name, index, default):
55        if name in kwargs:
56          return kwargs[name]
57        elif index < len(args):
58          return args[index]
59        else:
60          return default
61
62      if is_method:
63        name = "%s.%s" % (args[0].__class__.__name__, func.__name__)
64      else:
65        name = "%s.%s" % (func.__module__, func.__name__)
66
67      # Be sure to repr before calling func. Argument values may change.
68      arg_values = {
69          name: repr(get_arg_value(name, index, default))
70          for name, index, default in args_to_log}
71
72      start = trace_time.Now()
73      log.add_trace_event("B", start, category, name, arg_values)
74      try:
75        return func(*args, **kwargs)
76      finally:
77        end = trace_time.Now()
78        log.add_trace_event("E", end, category, name)
79    return traced_function
80
81  no_decorator_arguments = len(args) == 1 and callable(args[0])
82  if no_decorator_arguments:
83    arg_names = ()
84    return get_wrapper(args[0])
85  else:
86    arg_names = args
87    return get_wrapper
88