Issue #14697: Fix missing parser module support for set displays and set comprehensions.

mdickinson · mdickinson · commit 11c1dee183a3 · 2012-05-07T16:34:34.000+01:00
diff --git a/Lib/test/test_parser.py b/Lib/test/test_parser.py
@@ -305,6 +305,29 @@ def test_raise_statement(self):
                          "except Exception as e:\n"
                          "    raise ValueError from e\n")
 
+    def test_set_displays(self):
+        self.check_expr('{2}')
+        self.check_expr('{2,}')
+        self.check_expr('{2, 3}')
+        self.check_expr('{2, 3,}')
+
+    def test_dict_displays(self):
+        self.check_expr('{}')
+        self.check_expr('{a:b}')
+        self.check_expr('{a:b,}')
+        self.check_expr('{a:b, c:d}')
+        self.check_expr('{a:b, c:d,}')
+
+    def test_set_comprehensions(self):
+        self.check_expr('{x for x in seq}')
+        self.check_expr('{f(x) for x in seq}')
+        self.check_expr('{f(x) for x in seq if condition(x)}')
+
+    def test_dict_comprehensions(self):
+        self.check_expr('{x:x for x in seq}')
+        self.check_expr('{x**2:x[3] for x in seq if condition(x)}')
+        self.check_expr('{x:x for x in seq1 for y in seq2 if condition(x, y)}')
+
 
 #
 #  Second, we take *invalid* trees and make sure we get ParserError
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -61,6 +61,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #14697: Fix missing support for set displays and set comprehensions in
+  parser module.
+
 - Issue #14701: Fix missing support for 'raise ... from' in parser module.
 
 - Issue #13183: Fix pdb skipping frames after hitting a breakpoint and running
diff --git a/Modules/parsermodule.c b/Modules/parsermodule.c
@@ -2865,34 +2865,92 @@ validate_exprlist(node *tree)
                                     validate_expr_or_star_expr, "exprlist"));
 }
 
-
+/*
+ *  dictorsetmaker:
+ *
+ *  (test ':' test (comp_for | (',' test ':' test)* [','])) |
+ *  (test (comp_for | (',' test)* [',']))
+ */
 static int
 validate_dictorsetmaker(node *tree)
 {
     int nch = NCH(tree);
-    int res = (validate_ntype(tree, dictorsetmaker)
-               && (nch >= 3)
-               && validate_test(CHILD(tree, 0))
-               && validate_colon(CHILD(tree, 1))
-               && validate_test(CHILD(tree, 2)));
+    int res;
+    int i = 0;
+
+    res = validate_ntype(tree, dictorsetmaker);
+    if (!res)
+        return 0;
 
-    if (res && ((nch % 4) == 0))
-        res = validate_comma(CHILD(tree, --nch));
-    else if (res)
-        res = ((nch % 4) == 3);
-
-    if (res && (nch > 3)) {
-        int pos = 3;
-        /*  ( ',' test ':' test )*  */
-        while (res && (pos < nch)) {
-            res = (validate_comma(CHILD(tree, pos))
-                   && validate_test(CHILD(tree, pos + 1))
-                   && validate_colon(CHILD(tree, pos + 2))
-                   && validate_test(CHILD(tree, pos + 3)));
-            pos += 4;
+    if (nch - i < 1) {
+        (void) validate_numnodes(tree, 1, "dictorsetmaker");
+        return 0;
+    }
+
+    res = validate_test(CHILD(tree, i++));
+    if (!res)
+        return 0;
+
+    if (nch - i >= 2 && TYPE(CHILD(tree, i)) == COLON) {
+        /* Dictionary display or dictionary comprehension. */
+        res = (validate_colon(CHILD(tree, i++))
+               && validate_test(CHILD(tree, i++)));
+        if (!res)
+            return 0;
+
+        if (nch - i >= 1 && TYPE(CHILD(tree, i)) == comp_for) {
+            /* Dictionary comprehension. */
+            res = validate_comp_for(CHILD(tree, i++));
+            if (!res)
+                return 0;
+        }
+        else {
+            /* Dictionary display. */
+            while (nch - i >= 4) {
+                res = (validate_comma(CHILD(tree, i++))
+                       && validate_test(CHILD(tree, i++))
+                       && validate_colon(CHILD(tree, i++))
+                       && validate_test(CHILD(tree, i++)));
+                if (!res)
+                    return 0;
+            }
+            if (nch - i == 1) {
+                res = validate_comma(CHILD(tree, i++));
+                if (!res)
+                    return 0;
+            }
         }
     }
-    return (res);
+    else {
+        /* Set display or set comprehension. */
+        if (nch - i >= 1 && TYPE(CHILD(tree, i)) == comp_for) {
+            /* Set comprehension. */
+            res = validate_comp_for(CHILD(tree, i++));
+            if (!res)
+                return 0;
+        }
+        else {
+            /* Set display. */
+            while (nch - i >= 2) {
+                res = (validate_comma(CHILD(tree, i++))
+                       && validate_test(CHILD(tree, i++)));
+                if (!res)
+                    return 0;
+            }
+            if (nch - i == 1) {
+                res = validate_comma(CHILD(tree, i++));
+                if (!res)
+                    return 0;
+            }
+        }
+    }
+
+    if (nch - i > 0) {
+        err_string("Illegal trailing nodes for dictorsetmaker.");
+        return 0;
+    }
+
+    return 1;
 }
 
 
