Merge pull request #200 from dotjrich/list-extend

Adds support for LIST_EXTEND opcode

ASTree.cpp

@@ -1455,6 +1455,33 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 }
             }
             break;
+        case Pyc::LIST_EXTEND_A:
+            {
+                PycRef<ASTNode> rhs = stack.top();
+                stack.pop();
+                PycRef<ASTList> lhs = stack.top().cast<ASTList>();
+                stack.pop();
+
+                if (rhs.type() != ASTNode::NODE_OBJECT) {
+                    fprintf(stderr, "Unsupported argument found for LIST_EXTEND\n");
+                    break;
+                }
+
+                // I've only ever seen this be a SMALL_TUPLE, but let's be careful...
+                PycRef<PycObject> obj = rhs.cast<ASTObject>()->object();
+                if (obj->type() != PycObject::TYPE_TUPLE && obj->type() != PycObject::TYPE_SMALL_TUPLE) {
+                    fprintf(stderr, "Unsupported argument type found for LIST_EXTEND\n");
+                    break;
+                }
+
+                ASTList::value_t result = lhs->values();
+                for (const auto& it : obj.cast<PycTuple>()->values()) {
+                    result.push_back(new ASTObject(it));
+                }
+
+                stack.push(new ASTList(result));
+            }
+            break;
         case Pyc::LOAD_ATTR_A:
             {
                 PycRef<ASTNode> name = stack.top();

tests/input/list_extend.py

@@ -0,0 +1,9 @@
+a = [1, 2, 3]
+b = ['4', 5, 6]
+c = [7, (8, 9, 10)]
+d = [None]
+e = [('a', 'b', 'c')]
+f = [a, b]
+g = []
+h = [1, 2, [3, 4]]
+a = [None, None, None]

tests/tokenized/list_extend.txt

@@ -0,0 +1,9 @@
+a = [ 1 , 2 , 3 ] <EOL>
+b = [ '4' , 5 , 6 ] <EOL>
+c = [ 7 , ( 8 , 9 , 10 ) ] <EOL>
+d = [ None ] <EOL>
+e = [ ( 'a' , 'b' , 'c' ) ] <EOL>
+f = [ a , b ] <EOL>
+g = [ ] <EOL>
+h = [ 1 , 2 , [ 3 , 4 ] ] <EOL>
+a = [ None , None , None ] <EOL>