Parse cell & free vars for Py 3.11+

ASTree.cpp

@@ -2584,6 +2584,19 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 stack.push(value);
             }
             break;
+        case Pyc::MAKE_CELL_A:
+            // no op
+            break;
+        case Pyc::COPY_FREE_VARS_A:
+            // nonlocal should be inserted
+            // some nonlocals may be redundant but we're not optimizing currently
+            if (operand != code->freeVars()->size()) {
+                throw std::runtime_error("Different number of free variables copied");
+            }
+            for (int i = 0; i < code->freeVars()->size(); i++) {
+                code->markNonLocal(code->freeVars()->get(i).cast<PycString>());
+            }
+            break;
         default:
             fprintf(stderr, "Unsupported opcode: %s (%d)\n", Pyc::OpcodeName(opcode), opcode);
             cleanBuild = false;
@@ -3567,6 +3580,21 @@ void decompyle(PycRef<PycCode> code, PycModule* mod, std::ostream& pyc_output)
             }
             pyc_output << "\n";
         }
+
+        PycCode::nonlocals_t non_locals = code->getNonLocals();
+        if (non_locals.size()) {
+            start_line(cur_indent + 1, pyc_output);
+            pyc_output << "nonlocal ";
+            bool first = true;
+            for (const auto& non_local : non_locals) {
+                if (!first)
+                    pyc_output << ", ";
+                pyc_output << non_local->value();
+                first = false;
+            }
+            pyc_output << "\n";
+        }
+
         printDocstringAndGlobals = false;
     }
 

pyc_code.cpp

@@ -26,6 +26,7 @@ lntable                         Obj     Obj     Obj     Obj     Obj     Obj
 exceptiontable                                                          Obj
 */
 
+// TODO : Simplify this function
 void PycCode::load(PycData* stream, PycModule* mod)
 {
     if (mod->verCompare(1, 3) >= 0 && mod->verCompare(2, 3) < 0)
@@ -75,23 +76,49 @@ void PycCode::load(PycData* stream, PycModule* mod)
     m_consts = LoadObject(stream, mod).cast<PycSequence>();
     m_names = LoadObject(stream, mod).cast<PycSequence>();
 
+    // Represents localsplusnames for py 3.11+
     if (mod->verCompare(1, 3) >= 0)
         m_localNames = LoadObject(stream, mod).cast<PycSequence>();
     else
         m_localNames = new PycTuple;
 
-    if (mod->verCompare(3, 11) >= 0)
+    PycSimpleSequence::value_t free_vars_extra, cell_vars_extra;
+
+    if (mod->verCompare(3, 11) >= 0) {
         m_localKinds = LoadObject(stream, mod).cast<PycString>();
+
+        if (m_localKinds->length() != m_localNames->size()) {
+            throw std::runtime_error("All variables kinds are not available");
+        }
+
+        // Starting py3.11, all variables are now part of localsplusnames
+
+        for (int i = 0; i < m_localKinds->length(); i++) {
+            const char kind = m_localKinds->value()[i];
+            auto name = m_localNames->get(i);
+
+            if (kind & Kinds::CO_FAST_CELL) {
+                cell_vars_extra.push_back(name);
+            }
+            else if (kind & Kinds::CO_FAST_FREE) {
+                free_vars_extra.push_back(name);
+            }
+        }
+    }
     else
         m_localKinds = new PycString;
 
     if (mod->verCompare(2, 1) >= 0 && mod->verCompare(3, 11) < 0)
         m_freeVars = LoadObject(stream, mod).cast<PycSequence>();
+    else if (mod->verCompare(3, 11) >= 0)
+        m_freeVars = new PycTuple(free_vars_extra);
     else
         m_freeVars = new PycTuple;
 
     if (mod->verCompare(2, 1) >= 0 && mod->verCompare(3, 11) < 0)
         m_cellVars = LoadObject(stream, mod).cast<PycSequence>();
+    else if (mod->verCompare(3, 11) >= 0)
+        m_cellVars = new PycTuple(cell_vars_extra);
     else
         m_cellVars = new PycTuple;
 

pyc_code.h

@@ -22,7 +22,8 @@ class PycExceptionTableEntry {
 
 class PycCode : public PycObject {
 public:
-    typedef std::vector<PycRef<PycString>> globals_t;
+    typedef std::vector<PycRef<PycString>> globals_t, nonlocals_t;
+
     enum CodeFlags {
         CO_OPTIMIZED = 0x1,                                 // 1.3 ->
         CO_NEWLOCALS = 0x2,                                 // 1.3 ->
@@ -49,6 +50,17 @@ class PycCode : public PycObject {
         CO_NO_MONITORING_EVENTS = 0x2000000,                // 3.13 ->
     };
 
+    enum Kinds {                                            // 3.11 ->
+        CO_FAST_ARG_POS = 0x2,
+        CO_FAST_ARG_KW = 0x4,
+        CO_FAST_ARG_VAR = 0x8,
+        CO_FAST_ARG = CO_FAST_ARG_POS | CO_FAST_ARG_KW | CO_FAST_ARG_VAR ,
+        CO_FAST_HIDDEN = 0x10,
+        CO_FAST_LOCAL = 0x20,
+        CO_FAST_CELL = 0x40,
+        CO_FAST_FREE = 0x80,
+    };
+
     PycCode(int type = TYPE_CODE)
         : PycObject(type), m_argCount(), m_posOnlyArgCount(), m_kwOnlyArgCount(),
           m_numLocals(), m_stackSize(), m_flags(), m_firstLine() { }
@@ -99,6 +111,12 @@ class PycCode : public PycObject {
         m_globalsUsed.emplace_back(std::move(varname));
     }
 
+    const nonlocals_t& getNonLocals() const { return m_nonlocalsUsed; }
+
+    void markNonLocal(PycRef<PycString> varname) {
+        m_nonlocalsUsed.emplace_back(std::move(varname));
+    }
+
     std::vector<PycExceptionTableEntry> exceptionTableEntries() const;
 
 private:
@@ -118,6 +136,7 @@ class PycCode : public PycObject {
     PycRef<PycString> m_lnTable;
     PycRef<PycString> m_exceptTable;
     globals_t m_globalsUsed; /* Global vars used in this code */
+    nonlocals_t m_nonlocalsUsed; // Nonlocal vars used in this code
 };
 
 #endif

pyc_sequence.h

@@ -38,6 +38,11 @@ class PycTuple : public PycSimpleSequence {
     typedef PycSimpleSequence::value_t value_t;
     PycTuple(int type = TYPE_TUPLE) : PycSimpleSequence(type) { }
 
+    PycTuple(value_t values, int type = TYPE_TUPLE): PycSimpleSequence(type) {
+        m_values = values;
+        m_size = values.size();
+    }
+
     void load(class PycData* stream, class PycModule* mod) override;
 };
 

tests/input/non_local.py

@@ -0,0 +1,7 @@
+def outer():
+    x = 0
+    def inner():
+        nonlocal x
+        x += 1
+        print(x)
+    return inner

tests/run_tests.py

@@ -95,7 +95,7 @@ def run_test(test_file):
             if not compiled_files:
                 status_line += '\033[33mXFAIL ({})\033[0m\n'.format(xfails)
             else:
-                status_line += '\033[32mPASS ({})\033[33m + XFAIL ()\033[0m\n' \
+                status_line += '\033[32mPASS ({})\033[33m + XFAIL ({})\033[0m\n' \
                                .format(len(compiled_files), xfails)
         else:
             status_line += '\033[32mPASS ({})\033[0m\n'.format(len(compiled_files))

tests/tokenized/non_local.txt

@@ -0,0 +1,10 @@
+def outer ( ) : <EOL>
+<INDENT>
+x = 0 <EOL>
+def inner ( ) : <EOL>
+<INDENT>
+nonlocal x <EOL>
+x += 1 <EOL>
+print ( x ) <EOL>
+<OUTDENT>
+return inner <EOL>