python/google/appengine/ext/ndb/utils.py

   1 """Low-level utilities used internally by NDB.
   2
   3 These are not meant for use by code outside NDB.
   4 """
   5
   6 import functools
   7 import logging
   8 import os
   9 import sys
  10 import threading
  11
  12 __all__ = []
  13
  14 DEBUG = True  # Set to False for some speedups
  15
  16
  17 def logging_debug(*args):
  18   # NOTE: If you want to see debug messages, set the logging level
  19   # manually to logging.DEBUG - 1; or for tests use -v -v -v (see below).
  20   if DEBUG and logging.getLogger().level < logging.DEBUG:
  21     logging.debug(*args)
  22
  23
  24 def wrapping(wrapped):
  25   # A decorator to decorate a decorator's wrapper.  Following the lead
  26   # of Twisted and Monocle, this is supposed to make debugging heavily
  27   # decorated code easier.  We'll see...
  28   # TODO(pcostello): This copies the functionality of functools.wraps
  29   # following the patch in http://bugs.python.org/issue3445. We can replace
  30   # this once upgrading to python 3.3.
  31   def wrapping_wrapper(wrapper):
  32     try:
  33       wrapper.__wrapped__ = wrapped
  34       wrapper.__name__ = wrapped.__name__
  35       wrapper.__doc__ = wrapped.__doc__
  36       wrapper.__dict__.update(wrapped.__dict__)
  37       # Local functions won't have __module__ attribute.
  38       if hasattr(wrapped, '__module__'):
  39         wrapper.__module__ = wrapped.__module__
  40     except Exception:
  41       pass
  42     return wrapper
  43   return wrapping_wrapper
  44
  45
  46 # Define a base class for classes that need to be thread-local.
  47 # This is pretty subtle; we want to use threading.local if threading
  48 # is supported, but object if it is not.
  49 if threading.local.__module__ == 'thread':
  50   logging_debug('Using threading.local')
  51   threading_local = threading.local
  52 else:
  53   logging_debug('Not using threading.local')
  54   threading_local = object
  55
  56
  57 def get_stack(limit=10):
  58   # Return a list of strings showing where the current frame was called.
  59   if not DEBUG:
  60     return ()
  61   frame = sys._getframe(1)  # Always skip get_stack() itself.
  62   lines = []
  63   while len(lines) < limit and frame is not None:
  64     f_locals = frame.f_locals
  65     ndb_debug = f_locals.get('__ndb_debug__')
  66     if ndb_debug != 'SKIP':
  67       line = frame_info(frame)
  68       if ndb_debug is not None:
  69         line += ' # ' + str(ndb_debug)
  70       lines.append(line)
  71     frame = frame.f_back
  72   return lines
  73
  74
  75 def func_info(func, lineno=None):
  76   if not DEBUG:
  77     return None
  78   func = getattr(func, '__wrapped__', func)
  79   code = getattr(func, 'func_code', None)
  80   return code_info(code, lineno)
  81
  82
  83 def gen_info(gen):
  84   if not DEBUG:
  85     return None
  86   frame = gen.gi_frame
  87   if gen.gi_running:
  88     prefix = 'running generator '
  89   elif frame:
  90     if frame.f_lasti < 0:
  91       prefix = 'initial generator '
  92     else:
  93       prefix = 'suspended generator '
  94   else:
  95     prefix = 'terminated generator '
  96   if frame:
  97     return prefix + frame_info(frame)
  98   code = getattr(gen, 'gi_code', None)
  99   if code:
 100     return prefix + code_info(code)
 101   return prefix + hex(id(gen))
 102
 103
 104 def frame_info(frame):
 105   if not DEBUG:
 106     return None
 107   return code_info(frame.f_code, frame.f_lineno)
 108
 109
 110 def code_info(code, lineno=None):
 111   if not DEBUG or not code:
 112     return ''
 113   funcname = code.co_name
 114   # TODO: Be cleverer about stripping filename,
 115   # e.g. strip based on sys.path.
 116   filename = os.path.basename(code.co_filename)
 117   if lineno is None:
 118     lineno = code.co_firstlineno
 119   return '%s(%s:%s)' % (funcname, filename, lineno)
 120
 121
 122 def positional(max_pos_args):
 123   """A decorator to declare that only the first N arguments may be positional.
 124
 125   Note that for methods, n includes 'self'.
 126   """
 127   __ndb_debug__ = 'SKIP'
 128   def positional_decorator(wrapped):
 129     if not DEBUG:
 130       return wrapped
 131     __ndb_debug__ = 'SKIP'
 132     @wrapping(wrapped)
 133     def positional_wrapper(*args, **kwds):
 134       __ndb_debug__ = 'SKIP'
 135       if len(args) > max_pos_args:
 136         plural_s = ''
 137         if max_pos_args != 1:
 138           plural_s = 's'
 139         raise TypeError(
 140           '%s() takes at most %d positional argument%s (%d given)' %
 141           (wrapped.__name__, max_pos_args, plural_s, len(args)))
 142       return wrapped(*args, **kwds)
 143     return positional_wrapper
 144   return positional_decorator
 145
 146
 147 def decorator(wrapped_decorator):
 148   """Converts a function into a decorator that optionally accepts keyword
 149   arguments in its declaration.
 150
 151   Example usage:
 152     @utils.decorator
 153     def decorator(func, args, kwds, op1=None):
 154       ... apply op1 ...
 155       return func(*args, **kwds)
 156
 157     # Form (1), vanilla
 158     @decorator
 159     foo(...)
 160       ...
 161
 162     # Form (2), with options
 163     @decorator(op1=5)
 164     foo(...)
 165       ...
 166
 167   Args:
 168     wrapped_decorator: A function that accepts positional args (func, args,
 169       kwds) and any additional supported keyword arguments.
 170
 171   Returns:
 172     A decorator with an additional 'wrapped_decorator' property that is set to
 173   the original function.
 174   """
 175   def helper(_func=None, **options):
 176     def outer_wrapper(func):
 177       @wrapping(func)
 178       def inner_wrapper(*args, **kwds):
 179         return wrapped_decorator(func, args, kwds, **options)
 180       return inner_wrapper
 181
 182     if _func is None:
 183       # Form (2), with options.
 184       return outer_wrapper
 185
 186     # Form (1), vanilla.
 187     if options:
 188       # Don't allow @decorator(foo, op1=5).
 189       raise TypeError('positional arguments not supported')
 190     return outer_wrapper(_func)
 191   helper.wrapped_decorator = wrapped_decorator
 192   return helper
 193
 194
 195 def tweak_logging():
 196   # Hack for running tests with verbose logging.  If there are two or
 197   # more -v flags, turn on INFO logging; if there are 3 or more, DEBUG.
 198   # (A single -v just tells unittest.main() to print the name of each
 199   # test; we don't want to interfere with that.)
 200   # Also, if there is a -q flag, set DEBUG to False, suppressing more
 201   # debug info even from warnings.
 202   q = 0
 203   v = 0
 204   for arg in sys.argv[1:]:
 205     if arg.startswith('-v'):
 206       v += arg.count('v')
 207     if arg.startswith('-q'):
 208       q += arg.count('q')
 209   if v >= 2:
 210     level = logging.INFO
 211     if v >= 3:
 212       level = logging.DEBUG - 1
 213     logging.basicConfig(level=level)
 214   if q > 0:
 215     global DEBUG
 216     DEBUG = False
 217
 218
 219 if 'test' in os.path.basename(sys.argv[0]):
 220   tweak_logging()