nmigen/tracer.py

   1 import traceback
   2 import inspect
   3 from opcode import opname
   4
   5
   6 __all__ = ["NameNotFound", "get_var_name", "get_src_loc"]
   7
   8
   9 class NameNotFound(Exception):
  10     pass
  11
  12
  13 _raise_exception = object()
  14
  15
  16 def get_var_name(depth=2, default=_raise_exception):
  17     frame = inspect.currentframe()
  18     for _ in range(depth):
  19         frame = frame.f_back
  20
  21     code = frame.f_code
  22     call_index = frame.f_lasti
  23     while True:
  24         call_opc   = opname[code.co_code[call_index]]
  25         if call_opc in ("EXTENDED_ARG"):
  26             call_index += 2
  27         else:
  28             break
  29     if call_opc not in ("CALL_FUNCTION", "CALL_FUNCTION_KW", "CALL_FUNCTION_EX", "CALL_METHOD"):
  30         return None
  31
  32     index = call_index + 2
  33     while True:
  34         opc = opname[code.co_code[index]]
  35         if opc in ("STORE_NAME", "STORE_ATTR"):
  36             name_index = int(code.co_code[index + 1])
  37             return code.co_names[name_index]
  38         elif opc == "STORE_FAST":
  39             name_index = int(code.co_code[index + 1])
  40             return code.co_varnames[name_index]
  41         elif opc == "STORE_DEREF":
  42             name_index = int(code.co_code[index + 1])
  43             return code.co_cellvars[name_index]
  44         elif opc in ("LOAD_GLOBAL", "LOAD_ATTR", "LOAD_FAST", "LOAD_DEREF",
  45                      "DUP_TOP", "BUILD_LIST"):
  46             index += 2
  47         else:
  48             if default is _raise_exception:
  49                 raise NameNotFound
  50             else:
  51                 return default
  52
  53
  54 def get_src_loc(src_loc_at=0):
  55     # n-th  frame: get_src_loc()
  56     # n-1th frame: caller of get_src_loc() (usually constructor)
  57     # n-2th frame: caller of caller (usually user code)
  58     # Python returns the stack frames reversed, so it is enough to set limit and grab the very
  59     # first one in the array.
  60     tb = traceback.extract_stack(limit=3 + src_loc_at)
  61     return (tb[0].filename, tb[0].lineno)