Hide list comp variables and support set comprehensions

1 files changed, 209 insertions, 136 deletions

diff --git a/Python/symtable.c b/Python/symtable.c
index e9c9391..68deb0a 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -76,7 +76,7 @@ PySTEntry_New(struct symtable *st, identifier name, _Py_block_ty block,
 	ste->ste_generator = 0;
 	ste->ste_returns_value = 0;
 
-	if (PyDict_SetItem(st->st_symbols, ste->ste_id, (PyObject *)ste) < 0)
+	if (PyDict_SetItem(st->st_blocks, ste->ste_id, (PyObject *)ste) < 0)
 	    goto fail;
 	
 	return ste;
@@ -172,6 +172,8 @@ static int symtable_exit_block(struct symtable *st, void *ast);
 static int symtable_visit_stmt(struct symtable *st, stmt_ty s);
 static int symtable_visit_expr(struct symtable *st, expr_ty s);
 static int symtable_visit_genexp(struct symtable *st, expr_ty s);
+static int symtable_visit_listcomp(struct symtable *st, expr_ty s);
+static int symtable_visit_setcomp(struct symtable *st, expr_ty s);
 static int symtable_visit_arguments(struct symtable *st, arguments_ty);
 static int symtable_visit_excepthandler(struct symtable *st, excepthandler_ty);
 static int symtable_visit_alias(struct symtable *st, alias_ty);
@@ -186,7 +188,8 @@ static int symtable_implicit_arg(struct symtable *st, int pos);
 static int symtable_visit_annotations(struct symtable *st, stmt_ty s);
 
 
-static identifier top = NULL, lambda = NULL, genexpr = NULL;
+static identifier top = NULL, lambda = NULL, genexpr = NULL,
+    listcomp = NULL, setcomp = NULL;
 
 #define GET_IDENTIFIER(VAR) \
 	((VAR) ? (VAR) : ((VAR) = PyString_InternFromString(# VAR)))
@@ -204,14 +207,13 @@ symtable_new(void)
 		return NULL;
 
 	st->st_filename = NULL;
-	st->st_symbols = NULL;
+	st->st_blocks = NULL;
 
 	if ((st->st_stack = PyList_New(0)) == NULL)
 		goto fail;
-	if ((st->st_symbols = PyDict_New()) == NULL)
+	if ((st->st_blocks = PyDict_New()) == NULL)
 		goto fail; 
 	st->st_cur = NULL;
-	st->st_tmpname = 0;
 	st->st_private = NULL;
 	return st;
  fail:
@@ -230,6 +232,7 @@ PySymtable_Build(mod_ty mod, const char *filename, PyFutureFeatures *future)
 		return st;
 	st->st_filename = filename;
 	st->st_future = future;
+	/* Make the initial symbol information gathering pass */
 	if (!GET_IDENTIFIER(top) ||
 	    !symtable_enter_block(st, top, ModuleBlock, (void *)mod, 0)) {
 		PySymtable_Free(st);
@@ -238,7 +241,6 @@ PySymtable_Build(mod_ty mod, const char *filename, PyFutureFeatures *future)
 
 	st->st_top = st->st_cur;
 	st->st_cur->ste_unoptimized = OPT_TOPLEVEL;
-	/* Any other top-level initialization? */
 	switch (mod->kind) {
 	case Module_kind:
 		seq = mod->v.Module.body;
@@ -267,6 +269,7 @@ PySymtable_Build(mod_ty mod, const char *filename, PyFutureFeatures *future)
 		PySymtable_Free(st);
 		return NULL;
 	}
+	/* Make the second symbol analysis pass */
 	if (symtable_analyze(st))
 		return st;
 	PySymtable_Free(st);
@@ -280,7 +283,7 @@ PySymtable_Build(mod_ty mod, const char *filename, PyFutureFeatures *future)
 void
 PySymtable_Free(struct symtable *st)
 {
-	Py_XDECREF(st->st_symbols);
+	Py_XDECREF(st->st_blocks);
 	Py_XDECREF(st->st_stack);
 	PyMem_Free((void *)st);
 }
@@ -293,7 +296,7 @@ PySymtable_Lookup(struct symtable *st, void *key)
 	k = PyLong_FromVoidPtr(key);
 	if (k == NULL)
 		return NULL;
-	v = PyDict_GetItem(st->st_symbols, k);
+	v = PyDict_GetItem(st->st_blocks, k);
 	if (v) {
 		assert(PySTEntry_Check(v));
 		Py_INCREF(v);
@@ -314,7 +317,7 @@ PyST_GetScope(PySTEntryObject *ste, PyObject *name)
 	if (!v)
 		return 0;
 	assert(PyInt_Check(v));
-	return (PyInt_AS_LONG(v) >> SCOPE_OFF) & SCOPE_MASK;
+	return (PyInt_AS_LONG(v) >> SCOPE_OFFSET) & SCOPE_MASK;
 }
 
 
@@ -325,7 +328,7 @@ PyST_GetScope(PySTEntryObject *ste, PyObject *name)
    it determines which local variables are cell variables; they provide
    bindings that are used for free variables in enclosed blocks.  
 
-   There are also two kinds of free variables, implicit and explicit.  An 
+   There are also two kinds of global variables, implicit and explicit.  An 
    explicit global is declared with the global statement.  An implicit
    global is a free variable for which the compiler has found no binding
    in an enclosing function scope.  The implicit global is either a global
@@ -337,24 +340,30 @@ PyST_GetScope(PySTEntryObject *ste, PyObject *name)
    TODO(jhylton): Discuss nonlocal
 
    The symbol table requires two passes to determine the scope of each name.
-   The first pass collects raw facts from the AST: the name is a parameter 
-   here, the name is used by not defined here, etc.  The second pass analyzes
-   these facts during a pass over the PySTEntryObjects created during pass 1.
+   The first pass collects raw facts from the AST via the symtable_visit_*
+   functions: the name is a parameter here, the name is used but not defined
+   here, etc.  The second pass analyzes these facts during a pass over the
+   PySTEntryObjects created during pass 1.
 
    When a function is entered during the second pass, the parent passes
    the set of all name bindings visible to its children.  These bindings 
-   are used to determine if the variable is free or an implicit global.
+   are used to determine if non-local variables are free or implicit globals.
    After doing the local analysis, it analyzes each of its child blocks
-   using an updated set of name bindings.  
+   using an updated set of name bindings.
 
-   The children update the free variable set.  If a local variable is free 
-   in a child, the variable is marked as a cell.  The current function must 
-   provide runtime storage for the variable that may outlive the function's 
-   frame.  Cell variables are removed from the free set before the analyze
-   function returns to its parent.
+   The children update the free variable set.  If a local variable is added to
+   the free variable set by the child, the variable is marked as a cell.  The
+   function object being defined must provide runtime storage for the variable
+   that may outlive the function's frame.  Cell variables are removed from the
+   free set before the analyze function returns to its parent.
    
-   The sets of bound and free variables are implemented as dictionaries
-   mapping strings to None.
+   During analysis, the names are:
+      symbols: dict mapping from symbol names to flag values (including offset scope values)
+      scopes: dict mapping from symbol names to scope values (no offset)
+      local: set of all symbol names local to the current scope
+      bound: set of all symbol names local to a containing function scope
+      free: set of all symbol names referenced but not bound in child scopes
+      global: set of all symbol names explicitly declared as global
 */
 
 #define SET_SCOPE(DICT, NAME, I) { \
@@ -375,14 +384,14 @@ PyST_GetScope(PySTEntryObject *ste, PyObject *name)
 */
 
 static int 
-analyze_name(PySTEntryObject *ste, PyObject *dict, PyObject *name, long flags,
+analyze_name(PySTEntryObject *ste, PyObject *scopes, PyObject *name, long flags,
 	     PyObject *bound, PyObject *local, PyObject *free, 
 	     PyObject *global)
 {
 	if (flags & DEF_GLOBAL) {
 		if (flags & DEF_PARAM) {
 			PyErr_Format(PyExc_SyntaxError,
-				     "name '%s' is local and global",
+				     "name '%s' is parameter and global",
 				     PyString_AS_STRING(name));
 			return 0;
 		}
@@ -392,41 +401,37 @@ analyze_name(PySTEntryObject *ste, PyObject *dict, PyObject *name, long flags,
 				     PyString_AS_STRING(name));
 			return 0;
                 }
-		SET_SCOPE(dict, name, GLOBAL_EXPLICIT);
-		if (PyDict_SetItem(global, name, Py_None) < 0)
+		SET_SCOPE(scopes, name, GLOBAL_EXPLICIT);
+		if (PySet_Add(global, name) < 0)
+			return 0;
+		if (bound && (PySet_Discard(bound, name) < 0))
 			return 0;
-		if (bound && PyDict_GetItem(bound, name)) {
-			if (PyDict_DelItem(bound, name) < 0)
-				return 0;
-		}
 		return 1;
 	}
         if (flags & DEF_NONLOCAL) {
 		if (flags & DEF_PARAM) {
 			PyErr_Format(PyExc_SyntaxError,
-				     "name '%s' is local and nonlocal",
+				     "name '%s' is parameter and nonlocal",
 				     PyString_AS_STRING(name));
 			return 0;
 		}
-                if (!PyDict_GetItem(bound, name)) {
+                if (!PySet_Contains(bound, name)) {
                         PyErr_Format(PyExc_SyntaxError,
                                      "no binding for nonlocal '%s' found",
 				     PyString_AS_STRING(name));
                                      
                         return 0;
                 }
-                SET_SCOPE(dict, name, FREE);
+                SET_SCOPE(scopes, name, FREE);
                 ste->ste_free = 1;
-                return PyDict_SetItem(free, name, Py_None) >= 0;
+                return PySet_Add(free, name) >= 0;
         }
 	if (flags & DEF_BOUND) {
-		SET_SCOPE(dict, name, LOCAL);
-		if (PyDict_SetItem(local, name, Py_None) < 0)
+		SET_SCOPE(scopes, name, LOCAL);
+		if (PySet_Add(local, name) < 0)
+			return 0;
+		if (PySet_Discard(global, name) < 0)
 			return 0;
-		if (PyDict_GetItem(global, name)) {
-			if (PyDict_DelItem(global, name) < 0)
-				return 0;
-		}
 		return 1;
 	}
 	/* If an enclosing block has a binding for this name, it
@@ -434,21 +439,21 @@ analyze_name(PySTEntryObject *ste, PyObject *dict, PyObject *name, long flags,
 	   Note that having a non-NULL bound implies that the block
 	   is nested.
 	*/
-	if (bound && PyDict_GetItem(bound, name)) {
-		SET_SCOPE(dict, name, FREE);
+	if (bound && PySet_Contains(bound, name)) {
+		SET_SCOPE(scopes, name, FREE);
 		ste->ste_free = 1;
-		return PyDict_SetItem(free, name, Py_None) >= 0;
+		return PySet_Add(free, name) >= 0;
 	}
 	/* If a parent has a global statement, then call it global
 	   explicit?  It could also be global implicit.
 	 */
-	if (global && PyDict_GetItem(global, name)) {
-		SET_SCOPE(dict, name, GLOBAL_EXPLICIT);
+	if (global && PySet_Contains(global, name)) {
+		SET_SCOPE(scopes, name, GLOBAL_EXPLICIT);
 		return 1;
 	}
 	if (ste->ste_nested)
 		ste->ste_free = 1;
-	SET_SCOPE(dict, name, GLOBAL_IMPLICIT);
+	SET_SCOPE(scopes, name, GLOBAL_IMPLICIT);
 	return 1;
 }
 
@@ -463,35 +468,35 @@ analyze_name(PySTEntryObject *ste, PyObject *dict, PyObject *name, long flags,
 */
 
 static int
-analyze_cells(PyObject *scope, PyObject *free)
+analyze_cells(PyObject *scopes, PyObject *free)
 {
-        PyObject *name, *v, *w;
+        PyObject *name, *v, *v_cell;
 	int success = 0;
 	Py_ssize_t pos = 0;
 
-	w = PyInt_FromLong(CELL);
-	if (!w)
+	v_cell = PyInt_FromLong(CELL);
+	if (!v_cell)
 		return 0;
-	while (PyDict_Next(scope, &pos, &name, &v)) {
-		long flags;
+	while (PyDict_Next(scopes, &pos, &name, &v)) {
+		long scope;
 		assert(PyInt_Check(v));
-		flags = PyInt_AS_LONG(v);
-		if (flags != LOCAL)
+		scope = PyInt_AS_LONG(v);
+		if (scope != LOCAL)
 			continue;
-		if (!PyDict_GetItem(free, name))
+		if (!PySet_Contains(free, name))
 			continue;
 		/* Replace LOCAL with CELL for this name, and remove
 		   from free. It is safe to replace the value of name 
 		   in the dict, because it will not cause a resize.
 		 */
-		if (PyDict_SetItem(scope, name, w) < 0)
+		if (PyDict_SetItem(scopes, name, v_cell) < 0)
 			goto error;
-		if (!PyDict_DelItem(free, name) < 0)
+		if (PySet_Discard(free, name) < 0)
 			goto error;
 	}
 	success = 1;
  error:
-	Py_DECREF(w);
+	Py_DECREF(v_cell);
 	return success;
 }
 
@@ -526,77 +531,91 @@ check_unoptimized(const PySTEntryObject* ste) {
 	return 0;
 }
 
-/* Enter the final scope information into the st_symbols dict. 
+/* Enter the final scope information into the ste_symbols dict. 
  * 
  * All arguments are dicts.  Modifies symbols, others are read-only.
 */
 static int
-update_symbols(PyObject *symbols, PyObject *scope, 
+update_symbols(PyObject *symbols, PyObject *scopes, 
                PyObject *bound, PyObject *free, int classflag)
 {
-	PyObject *name, *v, *u, *w, *free_value = NULL;
+	PyObject *name = NULL, *itr = NULL;
+	PyObject *v = NULL, *v_scope = NULL, *v_new = NULL, *v_free = NULL;
 	Py_ssize_t pos = 0;
 
+	/* Update scope information for all symbols in this scope */
 	while (PyDict_Next(symbols, &pos, &name, &v)) {
-		long i, flags;
+		long scope, flags;
 		assert(PyInt_Check(v));
 		flags = PyInt_AS_LONG(v);
-		w = PyDict_GetItem(scope, name);
-		assert(w && PyInt_Check(w));
-		i = PyInt_AS_LONG(w);
-		flags |= (i << SCOPE_OFF);
-		u = PyInt_FromLong(flags);
-		if (!u)
+		v_scope = PyDict_GetItem(scopes, name);
+		assert(v_scope && PyInt_Check(v_scope));
+		scope = PyInt_AS_LONG(v_scope);
+		flags |= (scope << SCOPE_OFFSET);
+		v_new = PyInt_FromLong(flags);
+		if (!v_new)
 			return 0;
-		if (PyDict_SetItem(symbols, name, u) < 0) {
-			Py_DECREF(u);
+		if (PyDict_SetItem(symbols, name, v_new) < 0) {
+			Py_DECREF(v_new);
 			return 0;
 		}
-		Py_DECREF(u);
+		Py_DECREF(v_new);
 	}
 
-        free_value = PyInt_FromLong(FREE << SCOPE_OFF);
-        if (!free_value)
+	/* Record not yet resolved free variables from children (if any) */
+        v_free = PyInt_FromLong(FREE << SCOPE_OFFSET);
+        if (!v_free)
 		return 0;
 
-        /* add a free variable when it's only use is for creating a closure */
-        pos = 0;
-	while (PyDict_Next(free, &pos, &name, &v)) {
-		PyObject *o = PyDict_GetItem(symbols, name);
+	itr = PyObject_GetIter(free);
+	if (!itr)
+		goto error;
+
+	while ((name = PyIter_Next(itr))) {
+		v = PyDict_GetItem(symbols, name);
 
-		if (o) {
-			/* It could be a free variable in a method of
+		/* Handle symbol that already exists in this scope */
+		if (v) {
+			/* Handle a free variable in a method of
 			   the class that has the same name as a local
 			   or global in the class scope.
 			*/
 			if  (classflag && 
-			     PyInt_AS_LONG(o) & (DEF_BOUND | DEF_GLOBAL)) {
-				long i = PyInt_AS_LONG(o) | DEF_FREE_CLASS;
-				o = PyInt_FromLong(i);
-				if (!o) {
-					Py_DECREF(free_value);
-					return 0;
+			     PyInt_AS_LONG(v) & (DEF_BOUND | DEF_GLOBAL)) {
+				long flags = PyInt_AS_LONG(v) | DEF_FREE_CLASS;
+				v_new = PyInt_FromLong(flags);
+				if (!v_new) {
+					goto error;
 				}
-				if (PyDict_SetItem(symbols, name, o) < 0) {
-					Py_DECREF(o);
-					Py_DECREF(free_value);
-					return 0;
+				if (PyDict_SetItem(symbols, name, v_new) < 0) {
+					Py_DECREF(v_new);
+					goto error;
 				}
-				Py_DECREF(o);
+				Py_DECREF(v_new);
 			}
-			/* else it's not free, probably a cell */
+			/* It's a cell, or already a free variable in this scope */
+			Py_DECREF(name);
 			continue;
 		}
-		if (!PyDict_GetItem(bound, name))
+		/* Handle global symbol */
+		if (!PySet_Contains(bound, name)) {
+			Py_DECREF(name);
 			continue;       /* it's a global */
-
-		if (PyDict_SetItem(symbols, name, free_value) < 0) {
-			Py_DECREF(free_value);
-			return 0;
 		}
+		/* Propagate new free symbol up the lexical stack */
+		if (PyDict_SetItem(symbols, name, v_free) < 0) {
+			goto error;
+		}
+		Py_DECREF(name);
         }
-        Py_DECREF(free_value);
+	Py_DECREF(itr);
+        Py_DECREF(v_free);
 	return 1;
+error:
+	Py_XDECREF(v_free);
+	Py_XDECREF(itr);
+	Py_XDECREF(name);
+	return 0;
 }   
 
 /* Make final symbol table decisions for block of ste.
@@ -611,59 +630,74 @@ static int
 analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free, 
 	      PyObject *global)
 {
-	PyObject *name, *v, *local = NULL, *scope = NULL, *newbound = NULL;
+	PyObject *name, *v, *local = NULL, *scopes = NULL, *newbound = NULL;
 	PyObject *newglobal = NULL, *newfree = NULL;
 	int i, success = 0;
 	Py_ssize_t pos = 0;
 
-	local = PyDict_New();
-	if (!local)
+	scopes = PyDict_New();
+	if (!scopes)
 		goto error;
-	scope = PyDict_New();
-	if (!scope)
+	local = PySet_New(NULL);
+	if (!local)
 		goto error;
-	newglobal = PyDict_New();
+	newglobal = PySet_New(NULL);
 	if (!newglobal)
 		goto error;
-	newfree = PyDict_New();
+	newfree = PySet_New(NULL);
 	if (!newfree)
 		goto error;
-	newbound = PyDict_New();
+	newbound = PySet_New(NULL);
 	if (!newbound)
 		goto error;
 
+	/* Class namespace has no effect on names visible in
+	   nested functions, so populate the global and bound
+	   sets to be passed to child blocks before analyzing
+	   this one.
+	 */
 	if (ste->ste_type == ClassBlock) {
-		/* make a copy of globals before calling analyze_name(),
-		   because global statements in the class have no effect
-		   on nested functions.
-		*/
-		if (PyDict_Update(newglobal, global) < 0)
-			goto error;
-		if (bound)
-			if (PyDict_Update(newbound, bound) < 0)
+		/* Pass down previously bound symbols */
+		if (bound) {
+			if (!PyNumber_InPlaceOr(newbound, bound))
 				goto error;
+			Py_DECREF(newbound);
+		}
+		/* Pass down known globals */
+		if (!PyNumber_InPlaceOr(newglobal, global))
+			goto error;
+		Py_DECREF(newglobal);
 	}
 
+	/* Analyze symbols in current scope */
 	assert(PySTEntry_Check(ste));
 	assert(PyDict_Check(ste->ste_symbols));
 	while (PyDict_Next(ste->ste_symbols, &pos, &name, &v)) {
 		long flags = PyInt_AS_LONG(v);
-		if (!analyze_name(ste, scope, name, flags, bound, local, free,
+		if (!analyze_name(ste, scopes, name, flags, bound, local, free,
 				  global))
 			goto error;
 	}
 
+	/* Populate global and bound sets to be passed to children.
+	 */
 	if (ste->ste_type != ClassBlock) {
+		/* Add function locals to bound set */
 		if (ste->ste_type == FunctionBlock) {
-			if (PyDict_Update(newbound, local) < 0)
+			if (!PyNumber_InPlaceOr(newbound, local))
 				goto error;
+			Py_DECREF(newbound);
 		}
+		/* Pass down previously bound symbols */
 		if (bound) {
-			if (PyDict_Update(newbound, bound) < 0)
+			if (!PyNumber_InPlaceOr(newbound, bound))
 				goto error;
+			Py_DECREF(newbound);
 		}
-		if (PyDict_Update(newglobal, global) < 0)
+		/* Pass down known globals */
+		if (!PyNumber_InPlaceOr(newglobal, global))
 			goto error;
+		Py_DECREF(newglobal);
 	}
 
 	/* Recursively call analyze_block() on each child block */
@@ -674,24 +708,28 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
 		entry = (PySTEntryObject*)c;
 		if (!analyze_block(entry, newbound, newfree, newglobal))
 			goto error;
+		/* Check if any children have free variables */
 		if (entry->ste_free || entry->ste_child_free)
 			ste->ste_child_free = 1;
 	}
 
-	if (ste->ste_type == FunctionBlock && !analyze_cells(scope, newfree))
+	/* Check if any local variables need to be converted to cell variables */
+	if (ste->ste_type == FunctionBlock && !analyze_cells(scopes, newfree))
 		goto error;
-	if (!update_symbols(ste->ste_symbols, scope, bound, newfree,
+	/* Records the results of the analysis in the symbol table entry */
+	if (!update_symbols(ste->ste_symbols, scopes, bound, newfree,
 			    ste->ste_type == ClassBlock))
 		goto error;
 	if (!check_unoptimized(ste))
 		goto error;
 
-	if (PyDict_Update(free, newfree) < 0)
+	if (!PyNumber_InPlaceOr(free, newfree))
 		goto error;
+	Py_DECREF(free);
 	success = 1;
  error:
+	Py_XDECREF(scopes);
 	Py_XDECREF(local);
-	Py_XDECREF(scope);
 	Py_XDECREF(newbound);
 	Py_XDECREF(newglobal);
 	Py_XDECREF(newfree);
@@ -706,10 +744,10 @@ symtable_analyze(struct symtable *st)
 	PyObject *free, *global;
 	int r;
 
-	free = PyDict_New();
+	free = PySet_New(NULL);
 	if (!free)
 	    return 0;
-	global = PyDict_New();
+	global = PySet_New(NULL);
 	if (!global) {
 	    Py_DECREF(free);
 	    return 0;
@@ -1200,14 +1238,16 @@ symtable_visit_expr(struct symtable *st, expr_ty e)
 	case Set_kind:
 		VISIT_SEQ(st, expr, e->v.Set.elts);
 		break;
+        case GeneratorExp_kind:
+		if (!symtable_visit_genexp(st, e))
+			return 0;
+		break;
         case ListComp_kind:
-		if (!symtable_new_tmpname(st))
+		if (!symtable_visit_listcomp(st, e))
 			return 0;
-		VISIT(st, expr, e->v.ListComp.elt);
-		VISIT_SEQ(st, comprehension, e->v.ListComp.generators);
 		break;
-        case GeneratorExp_kind:
-		if (!symtable_visit_genexp(st, e))
+	case SetComp_kind:
+		if (!symtable_visit_setcomp(st, e))
 			return 0;
 		break;
         case Yield_kind:
@@ -1479,27 +1519,60 @@ symtable_visit_slice(struct symtable *st, slice_ty s)
 }
 
 static int 
-symtable_visit_genexp(struct symtable *st, expr_ty e)
+symtable_handle_comprehension(struct symtable *st, expr_ty e,
+                              identifier scope_name,
+                              asdl_seq *generators, expr_ty elt)
 {
+	int is_generator = (e->kind == GeneratorExp_kind);
+	int needs_tmp = !is_generator;
 	comprehension_ty outermost = ((comprehension_ty)
-			 (asdl_seq_GET(e->v.GeneratorExp.generators, 0)));
+                                        asdl_seq_GET(generators, 0));
 	/* Outermost iterator is evaluated in current scope */
 	VISIT(st, expr, outermost->iter);
-	/* Create generator scope for the rest */
-	if (!GET_IDENTIFIER(genexpr) ||
-	    !symtable_enter_block(st, genexpr, FunctionBlock, (void *)e, 0)) {
+	/* Create comprehension scope for the rest */
+	if (!scope_name ||
+	    !symtable_enter_block(st, scope_name, FunctionBlock, (void *)e, 0)) {
 		return 0;
 	}
-	st->st_cur->ste_generator = 1;
+	st->st_cur->ste_generator = is_generator;
 	/* Outermost iter is received as an argument */
 	if (!symtable_implicit_arg(st, 0)) {
 		symtable_exit_block(st, (void *)e);
 		return 0;
 	}
+	/* Allocate temporary name if needed */
+	if (needs_tmp && !symtable_new_tmpname(st)) {
+		symtable_exit_block(st, (void *)e);
+		return 0;
+	}
 	VISIT_IN_BLOCK(st, expr, outermost->target, (void*)e);
 	VISIT_SEQ_IN_BLOCK(st, expr, outermost->ifs, (void*)e);
 	VISIT_SEQ_TAIL_IN_BLOCK(st, comprehension,
-				e->v.GeneratorExp.generators, 1, (void*)e);
-	VISIT_IN_BLOCK(st, expr, e->v.GeneratorExp.elt, (void*)e);
+				generators, 1, (void*)e);
+	VISIT_IN_BLOCK(st, expr, elt, (void*)e);
 	return symtable_exit_block(st, (void *)e);
 }
+
+static int 
+symtable_visit_genexp(struct symtable *st, expr_ty e)
+{
+	return symtable_handle_comprehension(st, e, GET_IDENTIFIER(genexpr),
+                                             e->v.GeneratorExp.generators,
+                                             e->v.GeneratorExp.elt);
+}
+
+static int 
+symtable_visit_listcomp(struct symtable *st, expr_ty e)
+{
+	return symtable_handle_comprehension(st, e, GET_IDENTIFIER(listcomp),
+                                             e->v.ListComp.generators,
+                                             e->v.ListComp.elt);
+}
+
+static int
+symtable_visit_setcomp(struct symtable *st, expr_ty e)
+{
+	return symtable_handle_comprehension(st, e, GET_IDENTIFIER(setcomp),
+					     e->v.SetComp.generators,
+					     e->v.SetComp.elt);
+}