From 4959c766947c16387cbc729c47afe5135b8fc4bc Mon Sep 17 00:00:00 2001
From: rocky <rb@dustyfeet.com>
Date: Fri, 4 Nov 2022 02:05:34 -0400
Subject: [PATCH] More 3.0 list comprehension bug fixes

---
 uncompyle6/parsers/parse30.py     |  5 +++++
 uncompyle6/semantics/gencomp.py   | 18 +++++++++++++++---
 uncompyle6/semantics/n_actions.py |  6 +++++-
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/uncompyle6/parsers/parse30.py b/uncompyle6/parsers/parse30.py
index bf8815bf..93a47ebd 100644
--- a/uncompyle6/parsers/parse30.py
+++ b/uncompyle6/parsers/parse30.py
@@ -102,6 +102,11 @@ class Python30Parser(Python31Parser):
                              LOAD_FAST FOR_ITER store dict_comp_iter
                              JUMP_BACK _come_froms POP_TOP JUMP_BACK
 
+        dict_comp_func   ::= BUILD_MAP_0
+                             DUP_TOP STORE_FAST
+                             LOAD_ARG FOR_ITER store
+                             dict_comp_iter JUMP_BACK RETURN_VALUE RETURN_LAST
+
         stmt         ::= try_except30
         try_except30 ::= SETUP_EXCEPT suite_stmts_opt
                         _come_froms pt_bp
diff --git a/uncompyle6/semantics/gencomp.py b/uncompyle6/semantics/gencomp.py
index 4a3e4f00..b4295727 100644
--- a/uncompyle6/semantics/gencomp.py
+++ b/uncompyle6/semantics/gencomp.py
@@ -269,6 +269,7 @@ class ComprehensionMixin:
 
         is_30_dict_comp = False
         store = None
+
         if node == "list_comp_async":
             # We have two different kinds of grammar rules:
             #   list_comp_async ::= LOAD_LISTCOMP LOAD_STR MAKE_FUNCTION_0 expr ...
@@ -595,7 +596,7 @@ class ComprehensionMixin:
         collections = [node[-3]]
         list_ifs = []
 
-        if self.version[:2] == (3, 0) and n != "list_iter":
+        if self.version[:2] == (3, 0) and n.kind != "list_iter":
             # FIXME 3.0 is a snowflake here. We need
             # special code for this. Not sure if this is totally
             # correct.
@@ -638,7 +639,12 @@ class ComprehensionMixin:
                 # FIXME: adjust for set comprehension
                 if n == "list_for":
                     stores.append(n[2])
-                    n = n[3]
+                    if self.version[:2] == (3, 0):
+                        body_index = 5
+                    else:
+                        body_index = 3
+
+                    n = n[body_index]
                     if n[0] == "list_for":
                         # Dog-paddle down largely singleton reductions
                         # to find the collection (expr)
@@ -670,7 +676,11 @@ class ComprehensionMixin:
 
             assert n == "lc_body", tree
 
-            self.preorder(n[0])
+            if self.version[:2] == (3, 0):
+                body_index = 1
+            else:
+                body_index = 0
+            self.preorder(n[body_index])
 
         # FIXME: add indentation around "for"'s and "in"'s
         n_colls = len(collections)
@@ -684,6 +694,8 @@ class ComprehensionMixin:
                 self.write(" async")
                 pass
             self.write(" for ")
+            if self.version[:2] == (3, 0):
+                store = token
             self.preorder(store)
             self.write(" in ")
             self.preorder(collections[i])
diff --git a/uncompyle6/semantics/n_actions.py b/uncompyle6/semantics/n_actions.py
index 38569a61..1370f39b 100644
--- a/uncompyle6/semantics/n_actions.py
+++ b/uncompyle6/semantics/n_actions.py
@@ -1085,7 +1085,11 @@ class NonterminalActions:
     def n_set_comp(self, node):
         self.write("{")
         if node[0] in ["LOAD_SETCOMP", "LOAD_DICTCOMP"]:
-            self.comprehension_walk_newer(node, 1, 0)
+            if self.version == (3, 0):
+                iter_index = 6
+            else:
+                iter_index = 1
+            self.comprehension_walk_newer(node, iter_index=iter_index, code_index=0)
         elif node[0].kind == "load_closure" and self.version >= (3, 0):
             self.closure_walk(node, collection_index=4)
         else: