Python opcode.HAVE_ARGUMENT Examples

def _unpack_opargs(code, inserted_code_list, current_index):
    """
    Modified version of `_unpack_opargs` function from module `dis`.
    We have to use it, because sometimes code can be in an inconsistent state: if EXTENDED_ARG
    operator was introduced into the code, but it hasn't been inserted into `code_list` yet.
    In this case we can't use standard `_unpack_opargs` and we should check whether there are
    some new operators in `inserted_code_list`.
    """
    extended_arg = 0
    for i in range(0, len(code), 2):
        op = code[i]
        if op >= HAVE_ARGUMENT:
            if not extended_arg:
                # in case if we added EXTENDED_ARG, but haven't inserted it to the source code yet.
                for code_index in range(current_index, len(inserted_code_list)):
                    inserted_offset, inserted_code = inserted_code_list[code_index]
                    if inserted_offset == i and inserted_code[0] == EXTENDED_ARG:
                        extended_arg = inserted_code[1] << 8
            arg = code[i + 1] | extended_arg
            extended_arg = (arg << 8) if op == EXTENDED_ARG else 0
        else:
            arg = None
        yield (i, op, arg) 

def code_walker(code):
    l = len(code)
    code = array('B', code)
    i = 0
    extended_arg = 0
    while i < l:
        op = code[i]
        if op >= HAVE_ARGUMENT:
            oparg = code[i + 1] + code[i + 2] * 256 + extended_arg
            extended_arg = 0
            if op == EXTENDED_ARG:
                extended_arg = oparg * 65536
            yield i, (op, oparg)
            i += 3
        else:
            yield i, (op, None)
            i += 1 

def co_code_findloadednames(co):
    """Find in the code of a code object, all loaded names.
    (by LOAD_NAME, LOAD_GLOBAL or LOAD_FAST) """

    import dis
    from opcode import HAVE_ARGUMENT, opmap
    hasloadname = (opmap['LOAD_NAME'],
                   opmap['LOAD_GLOBAL'], opmap['LOAD_FAST'])
    insns = dis.get_instructions(co)
    len_co_names = len(co.co_names)
    indexset = {}
    for insn in insns:
        if insn.opcode >= HAVE_ARGUMENT:
            if insn.opcode in hasloadname:
                indexset[insn.argval] = 1
                if len(indexset) >= len_co_names:
                    break
    for name in co.co_varnames:
        try:
            del indexset[name]
        except KeyError:
            pass
    return indexset 

def decode_codeobj(codeobj):
    # adapted from dis.dis
    extended_arg = 0
    if is_py3k:
        codestr = codeobj.co_code
    else:
        codestr = [ord(ch) for ch in codeobj.co_code]
    free = None
    i = 0
    while i < len(codestr):
        op = codestr[i]
        opname = opcode.opname[op]
        i += 1
        argval = None
        if op >= opcode.HAVE_ARGUMENT:
            oparg = codestr[i] + codestr[i + 1] * 256 + extended_arg
            i += 2
            extended_arg = 0
            if op == opcode.EXTENDED_ARG:
                extended_arg = oparg * 65536
                continue
            
            if op in opcode.hasconst:
                argval = codeobj.co_consts[oparg]
            elif op in opcode.hasname:
                argval = codeobj.co_names[oparg]
            elif op in opcode.hasjrel:
                argval = i + oparg
            elif op in opcode.haslocal:
                argval = codeobj.co_varnames[oparg]
            elif op in opcode.hascompare:
                argval = opcode.cmp_op[oparg]
            elif op in opcode.hasfree:
                if free is None:
                    free = codeobj.co_cellvars + codeobj.co_freevars
                argval = free[oparg]

        yield (opname, argval) 

def visit(self, index, op, arg=None, lineno=None, cflow_in=False):
    """
      Callback of the visitor. It dynamically constructs the name
      of the specialized visitor to call based on the name of the opcode.

      :param index: Bytecode index.
      :param op: The opcode that is currently visited.
      :param arg: The expanded oparg (i.e., constants, names, etc. are resolved).
      :param lineno: The line number associated with the opcode.
      :param cflow_in: ``True`` if the current ``index`` is the target of a jump.
    """
    # Let's start with a slow impl of the jump table, with
    # reflection
    method_name = BytecodeVisitor.toMethodName(opcode.opname[op])
    if hasattr(self, method_name):
      meth = getattr(self, method_name)
      if op < opcode.HAVE_ARGUMENT:
        logger.debug("%03d %26s" % (lineno, method_name))
        return meth()
      else:
        logger.debug("%03d %26s( %s )" % (lineno, method_name, repr(arg)))
        return meth(arg)
    else:
      logger.error("Method not found: %s" % method_name)


  # 2.7 specific visitors. See https://docs.python.org/2/library/dis.html 

def instructions(code_obj):
    # easy for python 3.4+
    if sys.version_info >= (3, 4):
        for inst in dis.Bytecode(code_obj):
            yield inst
    else:
        # otherwise we have to manually parse
        code = code_obj.co_code
        NewInstruction = namedtuple('Instruction', ('opcode', 'arg'))
        if six.PY2:
            code = map(ord, code)
        i, L = 0, len(code)
        extended_arg = 0
        while i < L:
            op = code[i]
            i += 1
            if op < opcode.HAVE_ARGUMENT:
                yield NewInstruction(op, None)
                continue
            oparg = code[i] + (code[i + 1] << 8) + extended_arg
            extended_arg = 0
            i += 2
            if op == opcode.EXTENDED_ARG:
                extended_arg = oparg << 16
                continue
            yield NewInstruction(op, oparg) 

def disassemble(co):
    code = co.co_code
    labels = dis.findlabels(code)
    linestarts = dict(dis.findlinestarts(co))
    n = len(code)
    i = 0
    extended_arg = 0
    free = None
    lineno = None
    while i < n:
        c = code[i]
        op = ord(c)
        lineno = linestarts.get(i, lineno)
        is_label = i in labels
        ist = i
        i += 1
        if op >= opcode.HAVE_ARGUMENT:
            oparg = ord(code[i]) + ord(code[i + 1]) * 256 + extended_arg
            extended_arg = 0
            i += 2
            if op == opcode.EXTENDED_ARG:
                extended_arg = oparg * 65536L
            if op in opcode.hasconst:
                arg = co.co_consts[oparg]
            elif op in opcode.hasname:
                arg = co.co_names[oparg]
            elif op in opcode.hasjrel:
                arg = i + oparg
            elif op in opcode.haslocal:
                arg = co.co_varnames[oparg]
            elif op in opcode.hascompare:
                arg = opcode.cmp_op[oparg]
            elif op in opcode.hasfree:
                if free is None:
                    free = co.co_cellvars + co.co_freevars
                arg = free[oparg]
            else:
                arg = NOVAL
        else:
            arg = NOVAL
        yield ist, lineno, is_label, opcode.opname[op], arg 

def instructions(code_obj):
    # easy for python 3.4+
    if sys.version_info >= (3, 4):
        for inst in dis.Bytecode(code_obj):
            yield inst
    else:
        # otherwise we have to manually parse
        code = code_obj.co_code
        NewInstruction = namedtuple('Instruction', ('opcode', 'arg'))
        if six.PY2:
            code = map(ord, code)
        i, L = 0, len(code)
        extended_arg = 0
        while i < L:
            op = code[i]
            i+= 1
            if op < opcode.HAVE_ARGUMENT:
                yield NewInstruction(op, None)
                continue
            oparg = code[i] + (code[i+1] << 8) + extended_arg
            extended_arg = 0
            i += 2
            if op == opcode.EXTENDED_ARG:
                extended_arg = oparg << 16
                continue
            yield NewInstruction(op, oparg) 

def instructions(code_obj):
    # easy for python 3.4+
    if sys.version_info >= (3, 4):
        for inst in dis.Bytecode(code_obj):
            yield inst
    else:
        # otherwise we have to manually parse
        code = code_obj.co_code
        NewInstruction = namedtuple('Instruction', ('opcode', 'arg'))
        if six.PY2:
            code = map(ord, code)
        i, L = 0, len(code)
        extended_arg = 0
        while i < L:
            op = code[i]
            i += 1
            if op < opcode.HAVE_ARGUMENT:
                yield NewInstruction(op, None)
                continue
            oparg = code[i] + (code[i + 1] << 8) + extended_arg
            extended_arg = 0
            i += 2
            if op == opcode.EXTENDED_ARG:
                extended_arg = oparg << 16
                continue
            yield NewInstruction(op, oparg) 

def decode_codeobj(codeobj):
    # adapted from dis.dis
    extended_arg = 0
    if is_py3k:
        codestr = codeobj.co_code
    else:
        codestr = [ord(ch) for ch in codeobj.co_code]
    free = None
    i = 0
    while i < len(codestr):
        op = codestr[i]
        opname = opcode.opname[op]
        i += 1
        argval = None
        if op >= opcode.HAVE_ARGUMENT:
            oparg = codestr[i] + codestr[i + 1] * 256 + extended_arg
            i += 2
            extended_arg = 0
            if op == opcode.EXTENDED_ARG:
                extended_arg = oparg * 65536
                continue
            
            if op in opcode.hasconst:
                argval = codeobj.co_consts[oparg]
            elif op in opcode.hasname:
                argval = codeobj.co_names[oparg]
            elif op in opcode.hasjrel:
                argval = i + oparg
            elif op in opcode.haslocal:
                argval = codeobj.co_varnames[oparg]
            elif op in opcode.hascompare:
                argval = opcode.cmp_op[oparg]
            elif op in opcode.hasfree:
                if free is None:
                    free = codeobj.co_cellvars + codeobj.co_freevars
                argval = free[oparg]

        yield (opname, argval) 

def decode_codeobj(codeobj):
    # adapted from dis.dis
    extended_arg = 0
    if is_py3k:
        codestr = codeobj.co_code
    else:
        codestr = [ord(ch) for ch in codeobj.co_code]
    free = None
    i = 0
    while i < len(codestr):
        op = codestr[i]
        opname = opcode.opname[op]
        i += 1
        argval = None
        if op >= opcode.HAVE_ARGUMENT:
            oparg = codestr[i] + codestr[i + 1] * 256 + extended_arg
            i += 2
            extended_arg = 0
            if op == opcode.EXTENDED_ARG:
                extended_arg = oparg * 65536
                continue
            
            if op in opcode.hasconst:
                argval = codeobj.co_consts[oparg]
            elif op in opcode.hasname:
                argval = codeobj.co_names[oparg]
            elif op in opcode.hasjrel:
                argval = i + oparg
            elif op in opcode.haslocal:
                argval = codeobj.co_varnames[oparg]
            elif op in opcode.hascompare:
                argval = opcode.cmp_op[oparg]
            elif op in opcode.hasfree:
                if free is None:
                    free = codeobj.co_cellvars + codeobj.co_freevars
                argval = free[oparg]

        yield (opname, argval) 

def show_bytecode(bytecode, start=0, end=2**32):
  from ..analysis.python.effects import get_stack_effect

  if bytecode is None:
    return ''
  buffer = []
  j = start
  end = min(end, len(bytecode) - 1)
  while j <= end:
    index, lineno, op, arg, _, co = bytecode[j]
    uid = hex(id(co))[-5:]

    pop_push_str = ''
    try:
      pop, push = get_stack_effect(op, arg)
      pop_push_str = ' (-%d +%d) ' % (pop, push)
    except ValueError, ex:
      pop_push_str = '         '

    if op >= opcode.HAVE_ARGUMENT:
      rts = repr(arg)
      if len(rts) > 40:
        rts = rts[:40] + '[...]'
      jump_target = ''
      if op in opcode.hasjrel or op in opcode.hasjabs:
        jump_address = arg if op in opcode.hasjabs else index + arg + 3
        jump_target = ' -------------> (%4d)' % jump_address

      buffer.append("[%5s]%4d(%4d) %20s(%3d)%s (%s)%s"
                    % (uid, lineno, index, opcode.opname[op], op, pop_push_str, rts, jump_target))
    else:
      buffer.append("[%5s]%4d(%4d) %20s(%3d)%s"
                    % (uid, lineno, index, opcode.opname[op], op, pop_push_str))
    j += 1 

def get_instruction_size(self, op, arg=None, bc_index=0):
    op, oparg = self.get_op_oparg(op, arg, bc_index)
    if op < opcode.HAVE_ARGUMENT:
      return 1
    if oparg > 0xffff:
      return 5
    return 3 

def get_op_oparg(self, op, arg, bc_index=0):
    """
      Retrieve the opcode (`op`) and its argument (`oparg`) from the supplied
      opcode and argument.

      :param op: The current opcode.
      :param arg: The current dereferenced argument.
      :param bc_index: The current bytecode index.
    """

    # Conversion of LOAD_NAME opcode based on injected references (global, vars)
    if op == LOAD_NAME:
      if arg in self.name_to_fast:
        op = LOAD_FAST
      elif arg in self.name_to_global:
        op = LOAD_GLOBAL

    oparg = None
    if op >= opcode.HAVE_ARGUMENT:
      if op in opcode.hasconst:
        oparg = self.add_get_constant(arg)
      elif op in opcode.hasname:
        oparg = self.add_get_names(arg)
      elif op in opcode.haslocal:
        oparg = self.add_get_varnames(arg)
      elif op in opcode.hascompare:
        oparg = opcode.cmp_op.index(arg)
      elif op in opcode.hasfree:
        oparg = self.add_get_cellvars_freevars(arg)
      elif op in opcode.hasjrel or op in opcode.hasjabs:
        oparg = arg
      else:
        oparg = arg
    return op, oparg 

def hack_line_numbers(code):
    """ Replace a code object's line number information to claim that every
        byte of the bytecode is a new source line.  Returns a new code
        object.  Also recurses to hack the line numbers in nested code objects.
    """

    # Create a new lnotab table.  Each opcode is claimed to be at
    # 1000*lineno + (opcode number within line), so for example, the opcodes on
    # source line 12 will be given new line numbers 12000, 12001, 12002, etc.
    old_num = list(lnotab_numbers(code.co_lnotab, code.co_firstlineno))
    n_bytes = len(code.co_code)
    new_num = []
    line = 0
    opnum_in_line = 0
    i_byte = 0
    while i_byte < n_bytes:
        if old_num and i_byte == old_num[0][0]:
            line = old_num.pop(0)[1]
            opnum_in_line = 0
        new_num.append((i_byte, 100000000 + 1000*line + opnum_in_line))
        if ord(code.co_code[i_byte]) >= opcode.HAVE_ARGUMENT:
            i_byte += 3
        else:
            i_byte += 1
        opnum_in_line += 1

    # new_num is a list of pairs, (byteoff, lineoff).  Turn it into an lnotab.
    new_firstlineno = new_num[0][1]-1
    new_lnotab = lnotab_string(new_num, new_firstlineno)

    # Recurse into code constants in this code object.
    new_consts = []
    for const in code.co_consts:
        if type(const) == types.CodeType:
            new_consts.append(hack_line_numbers(const))
        else:
            new_consts.append(const)

    # Create a new code object, just like the old one, except with new
    # line numbers.
    new_code = new.code(
        code.co_argcount, code.co_nlocals, code.co_stacksize, code.co_flags,
        code.co_code, tuple(new_consts), code.co_names, code.co_varnames,
        code.co_filename, code.co_name, new_firstlineno, new_lnotab
        )

    return new_code 

def _update_label_offsets(code_obj, breakpoint_offset, breakpoint_code_list):
    """
    Update labels for the relative and absolute jump targets
    :param code_obj: code to modify
    :param breakpoint_offset: offset for the inserted code
    :param breakpoint_code_list: size of the inserted code
    :return: bytes sequence with modified labels; list of tuples (resulting offset, list of code instructions) with
    information about all inserted pieces of code
    """
    inserted_code = list()
    # the list with all inserted pieces of code
    inserted_code.append((breakpoint_offset, breakpoint_code_list))
    code_list = list(code_obj)
    j = 0

    while j < len(inserted_code):
        current_offset, current_code_list = inserted_code[j]
        offsets_for_modification = []

        for offset, op, arg in _unpack_opargs(code_list, inserted_code, j):
            if arg is not None:
                if op in dis.hasjrel:
                    # has relative jump target
                    label = offset + 2 + arg
                    if offset < current_offset < label:
                        # change labels for relative jump targets if code was inserted inside
                        offsets_for_modification.append(offset)
                elif op in dis.hasjabs:
                    # change label for absolute jump if code was inserted before it
                    if current_offset < arg:
                        offsets_for_modification.append(offset)
        for i in range(0, len(code_list), 2):
            op = code_list[i]
            if i in offsets_for_modification and op >= dis.HAVE_ARGUMENT:
                new_arg = code_list[i + 1] + len(current_code_list)
                if new_arg <= MAX_BYTE:
                    code_list[i + 1] = new_arg
                else:
                    # handle bytes overflow
                    if i - 2 > 0 and code_list[i - 2] == EXTENDED_ARG and code_list[i - 1] < MAX_BYTE:
                        # if new argument > 255 and EXTENDED_ARG already exists we need to increase it's argument
                        code_list[i - 1] += 1
                    else:
                        # if there isn't EXTENDED_ARG operator yet we have to insert the new operator
                        extended_arg_code = [EXTENDED_ARG, new_arg >> 8]
                        inserted_code.append((i, extended_arg_code))
                    code_list[i + 1] = new_arg & MAX_BYTE

        code_list = code_list[:current_offset] + current_code_list + code_list[current_offset:]

        for k in range(len(inserted_code)):
            offset, inserted_code_list = inserted_code[k]
            if current_offset < offset:
                inserted_code[k] = (offset + len(current_code_list), inserted_code_list)
        j += 1

    return bytes(code_list), inserted_code 

def hack_line_numbers(code):
    """ Replace a code object's line number information to claim that every
        byte of the bytecode is a new source line.  Returns a new code
        object.  Also recurses to hack the line numbers in nested code objects.
    """

    # Create a new lnotab table.  Each opcode is claimed to be at
    # 1000*lineno + (opcode number within line), so for example, the opcodes on
    # source line 12 will be given new line numbers 12000, 12001, 12002, etc.
    old_num = list(lnotab_numbers(code.co_lnotab, code.co_firstlineno))
    n_bytes = len(code.co_code)
    new_num = []
    line = 0
    opnum_in_line = 0
    i_byte = 0
    while i_byte < n_bytes:
        if old_num and i_byte == old_num[0][0]:
            line = old_num.pop(0)[1]
            opnum_in_line = 0
        new_num.append((i_byte, 100000000 + 1000*line + opnum_in_line))
        if ord(code.co_code[i_byte]) >= opcode.HAVE_ARGUMENT:
            i_byte += 3
        else:
            i_byte += 1
        opnum_in_line += 1

    # new_num is a list of pairs, (byteoff, lineoff).  Turn it into an lnotab.
    new_firstlineno = new_num[0][1]-1
    new_lnotab = lnotab_string(new_num, new_firstlineno)

    # Recurse into code constants in this code object.
    new_consts = []
    for const in code.co_consts:
        if type(const) == types.CodeType:
            new_consts.append(hack_line_numbers(const))
        else:
            new_consts.append(const)

    # Create a new code object, just like the old one, except with new
    # line numbers.
    new_code = new.code(
        code.co_argcount, code.co_nlocals, code.co_stacksize, code.co_flags,
        code.co_code, tuple(new_consts), code.co_names, code.co_varnames,
        code.co_filename, code.co_name, new_firstlineno, new_lnotab
        )

    return new_code 

def emit(self, op, oparg, arg=None, lineno=-1):
    """
      Writes the bytecode and lnotab.
    """
    bytecode_inc, line_inc = 0, -1

    if op >= opcode.HAVE_ARGUMENT and oparg > 0xffff:
      self.append_code(opcode.EXTENDED_ARG)
      self.append_code((oparg >> 16) & 0xff)
      self.append_code((oparg >> 24) & 0xff)
      bytecode_inc += 3

    self.append_code(op)
    bytecode_inc += 1

    if op >= opcode.HAVE_ARGUMENT:
      self.append_code(oparg & 0xff)
      self.append_code((oparg >> 8) & 0xff)
      bytecode_inc += 2

    # We also adjust the lnotabs field
    if self.prev_lineno == -1:
      line_inc = 0
    else:
      line_inc = lineno - self.prev_lineno

    if line_inc == 0 and bytecode_inc == 0:
      self.lnotab.append(0)  # bytecode increment
      self.lnotab.append(0)  # lineno increment
    else:
      while bytecode_inc > 255:
        self.lnotab.append(255)
        self.lnotab.append(0)
        bytecode_inc -= 255

      while line_inc > 255:
        self.lnotab.append(0)
        self.lnotab.append(255)
        line_inc -= 255

      # Add the remainder
      self.lnotab.append(bytecode_inc)
      self.lnotab.append(line_inc)

    self.prev_lineno = lineno