Hide list comp variables and support set comprehensions

4 files changed, 52 insertions, 44 deletions

diff --git a/Include/Python-ast.h b/Include/Python-ast.h
index e07f025..d88076b 100644
--- a/Include/Python-ast.h
+++ b/Include/Python-ast.h
@@ -183,10 +183,10 @@ struct _stmt {
 
 enum _expr_kind {BoolOp_kind=1, BinOp_kind=2, UnaryOp_kind=3, Lambda_kind=4,
                   IfExp_kind=5, Dict_kind=6, Set_kind=7, ListComp_kind=8,
-                  GeneratorExp_kind=9, Yield_kind=10, Compare_kind=11,
-                  Call_kind=12, Num_kind=13, Str_kind=14, Bytes_kind=15,
-                  Ellipsis_kind=16, Attribute_kind=17, Subscript_kind=18,
-                  Name_kind=19, List_kind=20, Tuple_kind=21};
+                  SetComp_kind=9, GeneratorExp_kind=10, Yield_kind=11,
+                  Compare_kind=12, Call_kind=13, Num_kind=14, Str_kind=15,
+                  Bytes_kind=16, Ellipsis_kind=17, Attribute_kind=18,
+                  Subscript_kind=19, Name_kind=20, List_kind=21, Tuple_kind=22};
 struct _expr {
         enum _expr_kind kind;
         union {
@@ -234,6 +234,11 @@ struct _expr {
                 struct {
                         expr_ty elt;
                         asdl_seq *generators;
+                } SetComp;
+                
+                struct {
+                        expr_ty elt;
+                        asdl_seq *generators;
                 } GeneratorExp;
                 
                 struct {
@@ -465,6 +470,9 @@ expr_ty _Py_Set(asdl_seq * elts, int lineno, int col_offset, PyArena *arena);
 #define ListComp(a0, a1, a2, a3, a4) _Py_ListComp(a0, a1, a2, a3, a4)
 expr_ty _Py_ListComp(expr_ty elt, asdl_seq * generators, int lineno, int
                      col_offset, PyArena *arena);
+#define SetComp(a0, a1, a2, a3, a4) _Py_SetComp(a0, a1, a2, a3, a4)
+expr_ty _Py_SetComp(expr_ty elt, asdl_seq * generators, int lineno, int
+                    col_offset, PyArena *arena);
 #define GeneratorExp(a0, a1, a2, a3, a4) _Py_GeneratorExp(a0, a1, a2, a3, a4)
 expr_ty _Py_GeneratorExp(expr_ty elt, asdl_seq * generators, int lineno, int
                          col_offset, PyArena *arena);
diff --git a/Include/graminit.h b/Include/graminit.h
index 79e40c4..0596fb1 100644
--- a/Include/graminit.h
+++ b/Include/graminit.h
@@ -46,10 +46,10 @@
 #define with_var 301
 #define except_clause 302
 #define suite 303
-#define testlist_safe 304
-#define old_test 305
-#define old_lambdef 306
-#define test 307
+#define test 304
+#define test_nocond 305
+#define lambdef 306
+#define lambdef_nocond 307
 #define or_test 308
 #define and_test 309
 #define not_test 310
@@ -64,25 +64,20 @@
 #define factor 319
 #define power 320
 #define atom 321
-#define listmaker 322
-#define testlist_gexp 323
-#define lambdef 324
-#define trailer 325
-#define subscriptlist 326
-#define subscript 327
-#define sliceop 328
-#define exprlist 329
-#define testlist 330
-#define dictsetmaker 331
-#define classdef 332
-#define arglist 333
-#define argument 334
-#define list_iter 335
-#define list_for 336
-#define list_if 337
-#define gen_iter 338
-#define gen_for 339
-#define gen_if 340
-#define testlist1 341
-#define encoding_decl 342
-#define yield_expr 343
+#define testlist_comp 322
+#define trailer 323
+#define subscriptlist 324
+#define subscript 325
+#define sliceop 326
+#define exprlist 327
+#define testlist 328
+#define dictorsetmaker 329
+#define classdef 330
+#define arglist 331
+#define argument 332
+#define comp_iter 333
+#define comp_for 334
+#define comp_if 335
+#define testlist1 336
+#define encoding_decl 337
+#define yield_expr 338
diff --git a/Include/opcode.h b/Include/opcode.h
index 662fbb4..5fbb037 100644
--- a/Include/opcode.h
+++ b/Include/opcode.h
@@ -21,6 +21,7 @@ extern "C" {
 
 #define UNARY_INVERT	15
 
+#define SET_ADD	17
 #define LIST_APPEND	18
 #define BINARY_POWER	19
 
diff --git a/Include/symtable.h b/Include/symtable.h
index 5f50105..da11603 100644
--- a/Include/symtable.h
+++ b/Include/symtable.h
@@ -5,31 +5,35 @@
 extern "C" {
 #endif
 
+/* XXX(ncoghlan): This is a weird mix of public names and interpreter internal
+ *                names.
+ */
+
 typedef enum _block_type { FunctionBlock, ClassBlock, ModuleBlock }
     _Py_block_ty;
 
 struct _symtable_entry;
 
 struct symtable {
-	const char *st_filename; /* name of file being compiled */
+	const char *st_filename;        /* name of file being compiled */
 	struct _symtable_entry *st_cur; /* current symbol table entry */
-	struct _symtable_entry *st_top; /* module entry */
-	PyObject *st_symbols;    /* dictionary of symbol table entries */
-        PyObject *st_stack;      /* stack of namespace info */
-	PyObject *st_global;     /* borrowed ref to MODULE in st_symbols */
-	int st_nblocks;          /* number of blocks */
-	PyObject *st_private;        /* name of current class or NULL */
-        int st_tmpname;          /* temporary name counter */
-	PyFutureFeatures *st_future; /* module's future features */
+	struct _symtable_entry *st_top; /* symbol table entry for module */
+	PyObject *st_blocks;            /* dict: map AST node addresses
+	                                 *       to symbol table entries */
+	PyObject *st_stack;             /* list: stack of namespace info */
+	PyObject *st_global;            /* borrowed ref to st_top->st_symbols */
+	int st_nblocks;                 /* number of blocks used */
+	PyObject *st_private;           /* name of current class or NULL */
+	PyFutureFeatures *st_future;    /* module's future features */
 };
 
 typedef struct _symtable_entry {
 	PyObject_HEAD
-	PyObject *ste_id;        /* int: key in st_symbols */
-	PyObject *ste_symbols;   /* dict: name to flags */
-	PyObject *ste_name;      /* string: name of block */
+	PyObject *ste_id;        /* int: key in ste_table->st_blocks */
+	PyObject *ste_symbols;   /* dict: variable names to flags */
+	PyObject *ste_name;      /* string: name of current block */
 	PyObject *ste_varnames;  /* list of variable names */
-	PyObject *ste_children;  /* list of child ids */
+	PyObject *ste_children;  /* list of child blocks */
 	_Py_block_ty ste_type;   /* module, class, or function */
 	int ste_unoptimized;     /* false if namespace is optimized */
 	unsigned ste_nested : 1;      /* true if block is nested */
@@ -80,7 +84,7 @@ PyAPI_FUNC(void) PySymtable_Free(struct symtable *);
    table.  GLOBAL is returned from PyST_GetScope() for either of them. 
    It is stored in ste_symbols at bits 12-15.
 */
-#define SCOPE_OFF 11
+#define SCOPE_OFFSET 11
 #define SCOPE_MASK (DEF_GLOBAL | DEF_LOCAL | DEF_PARAM | DEF_NONLOCAL)
 
 #define LOCAL 1