Handle extended args.

Author: ericsnowcurrently

Objects/codeobject.c

@@ -1690,7 +1690,27 @@ PyCode_GetFreevars(PyCodeObject *code)
 }
 
 
-#define GET_OPARG(co, i, initial) (initial)
+static _Py_CODEUNIT
+advance_past_extended_arg(PyCodeObject *co, int *p_i, int *p_oparg)
+{
+    int i = *p_i;
+    _Py_CODEUNIT inst = _Py_GetBaseCodeUnit(co, i);
+    int opcode = inst.op.code;
+    int oparg = 0;
+    while (opcode == EXTENDED_ARG) {
+        oparg = (oparg << 8) | inst.op.arg;
+        i++;
+        inst = _Py_GetBaseCodeUnit(co, i);
+        opcode = inst.op.code;
+    }
+    if (p_oparg != NULL) {
+        *p_oparg = (oparg << 8) | inst.op.arg;
+    }
+    *p_i = i;
+    return inst;
+}
+
+
 // We may want to move these macros to pycore_opcode_utils.h
 // and use them in Python/bytecodes.c.
 #define LOAD_GLOBAL_NAME_INDEX(oparg) ((oparg)>>1)
@@ -1731,9 +1751,9 @@ identify_unbound_names(PyThreadState *tstate, PyCodeObject *co,
     struct co_unbound_counts unbound = {0};
     Py_ssize_t len = Py_SIZE(co);
     for (int i = 0; i < len; i += _PyInstruction_GetLength(co, i)) {
-        _Py_CODEUNIT inst = _Py_GetBaseCodeUnit(co, i);
+        int oparg;
+        _Py_CODEUNIT inst = advance_past_extended_arg(co, &i, &oparg);
         if (inst.op.code == LOAD_ATTR) {
-            int oparg = GET_OPARG(co, i, inst.op.arg);
             int index = LOAD_ATTR_NAME_INDEX(oparg);
             PyObject *name = GETITEM(co->co_names, index);
             if (PySet_Contains(attrnames, name)) {
@@ -1749,7 +1769,6 @@ identify_unbound_names(PyThreadState *tstate, PyCodeObject *co,
             }
         }
         else if (inst.op.code == LOAD_GLOBAL) {
-            int oparg = GET_OPARG(co, i, inst.op.arg);
             int index = LOAD_ATTR_NAME_INDEX(oparg);
             PyObject *name = GETITEM(co->co_names, index);
             if (PySet_Contains(globalnames, name)) {
@@ -1983,7 +2002,8 @@ code_returns_only_none(PyCodeObject *co)
         if (IS_RETURN_OPCODE(inst.op.code)) {
             assert(i != 0);
             // Ignore it if it returns None.
-            _Py_CODEUNIT prev = _Py_GetBaseCodeUnit(co, i-1);
+            int prev_i = i - 1;
+            _Py_CODEUNIT prev = _Py_GetBaseCodeUnit(co, prev_i);
             if (prev.op.code == LOAD_CONST) {
                 // We don't worry about EXTENDED_ARG for now.
                 if (prev.op.arg == none_index) {