Destruct

Author: tomasr8

Grammar/python.gram

@@ -1135,6 +1135,20 @@ star_atom[expr_ty]:
     | '(' a=target_with_star_atom ')' { _PyPegen_set_expr_context(p, a, Store) }
     | '(' a=[star_targets_tuple_seq] ')' { _PyAST_Tuple(a, Store, EXTRA) }
     | '[' a=[star_targets_list_seq] ']' { _PyAST_List(a, Store, EXTRA) }
+    | '{' a=dict_unpack_fields '}' { a }
+
+dict_unpack_fields[expr_ty]:
+    | a[asdl_seq*]=','.dict_unpack_field+ ',' '**' b=NAME [','] {
+        _PyPegen_make_dict_unpack_rest(p, a, b, EXTRA) }
+    | '**' b=NAME [','] {
+        _PyPegen_make_dict_unpack_rest(p, NULL, b, EXTRA) }
+    | a[asdl_seq*]=','.dict_unpack_field+ [','] {
+        _PyPegen_make_dict_unpack(p, a, EXTRA) }
+
+dict_unpack_field[KeyValuePair*]:
+    | a=STRING ':' b=star_atom { _PyPegen_dict_unpack_kv(p, _PyPegen_constant_from_string(p, a), b) }
+    | a=NAME ':' b=star_atom { _PyPegen_dict_unpack_kv_name(p, a, b) }
+    | a=NAME { _PyPegen_dict_unpack_shorthand(p, a) }
 
 single_target[expr_ty]:
     | single_subscript_attribute_target

Include/internal/pycore_ast.h

@@ -29,8 +29,10 @@
 #define INTRINSIC_TYPEVAR_WITH_CONSTRAINTS       3
 #define INTRINSIC_SET_FUNCTION_TYPE_PARAMS       4
 #define INTRINSIC_SET_TYPEPARAM_DEFAULT          5
+#define INTRINSIC_DESTRUCTURE                    6
+#define INTRINSIC_DESTRUCTURE_REST               7
 
-#define MAX_INTRINSIC_2                          5
+#define MAX_INTRINSIC_2                          7
 
 typedef PyObject *(*intrinsic_func1)(PyThreadState* tstate, PyObject *value);
 typedef PyObject *(*intrinsic_func2)(PyThreadState* tstate, PyObject *value1, PyObject *value2);

Include/internal/pycore_ast_state.h

@@ -90,6 +90,7 @@ module Python
          | Name(identifier id, expr_context ctx)
          | List(expr* elts, expr_context ctx)
          | Tuple(expr* elts, expr_context ctx)
+         | DictUnpack(expr* keys, expr* targets, expr? rest, expr_context ctx)
 
          -- can appear only in Subscript
          | Slice(expr? lower, expr? upper, expr? step)

Include/internal/pycore_intrinsics.h

@@ -330,6 +330,14 @@ _PyPegen_set_expr_context(Parser *p, expr_ty expr, expr_context_ty ctx)
         case Starred_kind:
             new = _set_starred_context(p, expr, ctx);
             break;
+        case DictUnpack_kind:
+            new = _PyAST_DictUnpack(
+                    expr->v.DictUnpack.keys,
+                    _set_seq_context(p, expr->v.DictUnpack.targets, ctx),
+                    expr->v.DictUnpack.rest ? _PyPegen_set_expr_context(p, expr->v.DictUnpack.rest, ctx) : NULL,
+                    ctx,
+                    EXTRA_EXPR(expr, expr));
+            break;
         default:
             new = expr;
     }
@@ -381,6 +389,117 @@ _PyPegen_get_values(Parser *p, asdl_seq *seq)
     return new_seq;
 }
 
+/* Dict destructuring helpers */
+
+/* Shorthand: {name} -> key=Constant("name"), target=Name("name", Store) */
+KeyValuePair *
+_PyPegen_dict_unpack_shorthand(Parser *p, expr_ty name)
+{
+    assert(name->kind == Name_kind);
+    KeyValuePair *a = _PyArena_Malloc(p->arena, sizeof(KeyValuePair));
+    if (!a) {
+        return NULL;
+    }
+    a->key = _PyAST_Constant(name->v.Name.id, NULL, EXTRA_EXPR(name, name));
+    if (!a->key) {
+        return NULL;
+    }
+    a->value = _PyAST_Name(name->v.Name.id, Store, EXTRA_EXPR(name, name));
+    if (!a->value) {
+        return NULL;
+    }
+    return a;
+}
+
+/* Explicit string key: {'key': target} */
+KeyValuePair *
+_PyPegen_dict_unpack_kv(Parser *p, expr_ty string_expr, expr_ty target)
+{
+    KeyValuePair *a = _PyArena_Malloc(p->arena, sizeof(KeyValuePair));
+    if (!a) {
+        return NULL;
+    }
+    /* string_expr is a Constant from the STRING token */
+    a->key = string_expr;
+    a->value = target;
+    return a;
+}
+
+/* Name as key: {name: target} -> key=Constant("name"), target=target */
+KeyValuePair *
+_PyPegen_dict_unpack_kv_name(Parser *p, expr_ty name, expr_ty target)
+{
+    assert(name->kind == Name_kind);
+    KeyValuePair *a = _PyArena_Malloc(p->arena, sizeof(KeyValuePair));
+    if (!a) {
+        return NULL;
+    }
+    a->key = _PyAST_Constant(name->v.Name.id, NULL, EXTRA_EXPR(name, name));
+    if (!a->key) {
+        return NULL;
+    }
+    a->value = target;
+    return a;
+}
+
+/* Build a DictUnpack AST node from a sequence of KeyValuePair* */
+expr_ty
+_PyPegen_make_dict_unpack(Parser *p, asdl_seq *seq,
+                          int lineno, int col_offset,
+                          int end_lineno, int end_col_offset,
+                          PyArena *arena)
+{
+    Py_ssize_t len = asdl_seq_LEN(seq);
+    asdl_expr_seq *keys = _Py_asdl_expr_seq_new(len, arena);
+    if (!keys) {
+        return NULL;
+    }
+    asdl_expr_seq *targets = _Py_asdl_expr_seq_new(len, arena);
+    if (!targets) {
+        return NULL;
+    }
+    for (Py_ssize_t i = 0; i < len; i++) {
+        KeyValuePair *pair = asdl_seq_GET_UNTYPED(seq, i);
+        asdl_seq_SET(keys, i, pair->key);
+        asdl_seq_SET(targets, i, pair->value);
+    }
+    return _PyAST_DictUnpack(keys, targets, NULL, Store,
+                             lineno, col_offset, end_lineno, end_col_offset,
+                             arena);
+}
+
+/* Build a DictUnpack AST node with **rest from a sequence of KeyValuePair* */
+expr_ty
+_PyPegen_make_dict_unpack_rest(Parser *p, asdl_seq *seq, expr_ty rest_name,
+                               int lineno, int col_offset,
+                               int end_lineno, int end_col_offset,
+                               PyArena *arena)
+{
+    Py_ssize_t len = seq ? asdl_seq_LEN(seq) : 0;
+    asdl_expr_seq *keys = _Py_asdl_expr_seq_new(len, arena);
+    if (!keys) {
+        return NULL;
+    }
+    asdl_expr_seq *targets = _Py_asdl_expr_seq_new(len, arena);
+    if (!targets) {
+        return NULL;
+    }
+    for (Py_ssize_t i = 0; i < len; i++) {
+        KeyValuePair *pair = asdl_seq_GET_UNTYPED(seq, i);
+        asdl_seq_SET(keys, i, pair->key);
+        asdl_seq_SET(targets, i, pair->value);
+    }
+    /* rest_name is a Name expr in Load context; convert to Store */
+    expr_ty rest = _PyAST_Name(rest_name->v.Name.id, Store,
+                               EXTRA_EXPR(rest_name, rest_name));
+    if (!rest) {
+        return NULL;
+    }
+    return _PyAST_DictUnpack(keys, targets, rest, Store,
+                             lineno, col_offset, end_lineno, end_col_offset,
+                             arena);
+}
+
 /* Constructs a KeyPatternPair that is used when parsing mapping & class patterns */
 KeyPatternPair *
 _PyPegen_key_pattern_pair(Parser *p, expr_ty key, pattern_ty pattern)
@@ -1083,6 +1202,8 @@ _PyPegen_get_expr_name(expr_ty e)
             return "list";
         case Tuple_kind:
             return "tuple";
+        case DictUnpack_kind:
+            return "dict destructuring";
         case Lambda_kind:
             return "lambda";
         case Call_kind:
@@ -1221,6 +1342,17 @@ _PyPegen_get_invalid_target(expr_ty e, TARGETS_TYPE targets_type)
         case Tuple_kind:
             VISIT_CONTAINER(e, Tuple);
             return NULL;
+        case DictUnpack_kind: {
+            Py_ssize_t len = asdl_seq_LEN(e->v.DictUnpack.targets);
+            for (Py_ssize_t i = 0; i < len; i++) {
+                expr_ty other = asdl_seq_GET(e->v.DictUnpack.targets, i);
+                expr_ty child = _PyPegen_get_invalid_target(other, targets_type);
+                if (child != NULL) {
+                    return child;
+                }
+            }
+            return NULL;
+        }
         case Starred_kind:
             if (targets_type == DEL_TARGETS) {
                 return e;