Add dict unpacking (#1152)

slozier · web-flow · commit 2c684b33ae67 · 2021-03-28T20:03:15.000-04:00
* Add nullable annotations to DictionaryExpression

* Add dict unpacking
diff --git a/Src/IronPython/Compiler/Ast/AstMethods.cs b/Src/IronPython/Compiler/Ast/AstMethods.cs
@@ -77,7 +77,10 @@ internal static class AstMethods {
         public static readonly MethodInfo PushFrame = GetMethod((Func<CodeContext, FunctionCode, List<FunctionStack>>)PythonOps.PushFrame);
         public static readonly MethodInfo FormatString = GetMethod((Func<CodeContext, string, object, string>)PythonOps.FormatString);
         public static readonly MethodInfo GeneratorCheckThrowableAndReturnSendValue = GetMethod((Func<object, object>)PythonOps.GeneratorCheckThrowableAndReturnSendValue);
-        
+
+        // methods matching Python opcodes
+        public static readonly MethodInfo DictUpdate = GetMethod((Action<CodeContext, PythonDictionary, object>)PythonOps.DictUpdate);
+
         private static MethodInfo GetMethod(Delegate x) {
             return x.Method;
         }
diff --git a/Src/IronPython/Compiler/Ast/DictionaryExpression.cs b/Src/IronPython/Compiler/Ast/DictionaryExpression.cs
@@ -2,53 +2,101 @@
 // The .NET Foundation licenses this file to you under the Apache 2.0 License.
 // See the LICENSE file in the project root for more information.
 
-using MSAst = System.Linq.Expressions;
+#nullable enable
 
 using System;
 using System.Collections.Generic;
-
-using Microsoft.Scripting.Interpreter;
+using System.Diagnostics;
 
 using IronPython.Runtime;
 using IronPython.Runtime.Operations;
 
-using AstUtils = Microsoft.Scripting.Ast.Utils;
+using Microsoft.Scripting.Interpreter;
+
+using MSAst = System.Linq.Expressions;
 
 namespace IronPython.Compiler.Ast {
-    using Ast = MSAst.Expression;
-
     public class DictionaryExpression : Expression, IInstructionProvider {
         private readonly SliceExpression[] _items;
-        private static readonly MSAst.Expression EmptyDictExpression = Ast.Call(AstMethods.MakeEmptyDict);
+        private readonly bool _hasNullKey;
+        private static readonly MSAst.Expression EmptyDictExpression = Expression.Call(AstMethods.MakeEmptyDict);
 
         public DictionaryExpression(params SliceExpression[] items) {
+            foreach (var item in items) {
+                if (item.SliceStart is null) _hasNullKey = true;
+                if (item.SliceStop is null) throw PythonOps.ValueError("None disallowed in expression list");
+            }
             _items = items;
         }
 
-        public IList<SliceExpression> Items => _items;
+        public IReadOnlyList<SliceExpression> Items => _items;
 
         public override MSAst.Expression Reduce() {
+            // empty dictionary
+            if (_items.Length == 0) {
+                return EmptyDictExpression;
+            }
+
+            if (_hasNullKey) {
+                // TODO: unpack constant dicts?
+                return ReduceDictionaryWithUnpack(Parent.LocalContext, _items.AsSpan());
+            }
+
             // create keys & values into array and then call helper function
             // which creates the dictionary
-            if (_items.Length != 0) {
-                return ReduceConstant() ?? ReduceDictionaryWithItems();
-            }
+            return ReduceConstant() ?? ReduceDictionaryWithItems(_items.AsSpan());
+        }
 
-            // empty dictionary
-            return EmptyDictExpression;
+        private static MSAst.Expression ReduceDictionaryWithUnpack(MSAst.Expression context, ReadOnlySpan<SliceExpression> items) {
+            Debug.Assert(items.Length > 0);
+            var expressions = new List<MSAst.Expression>(items.Length + 2);
+            var varExpr = Expression.Variable(typeof(PythonDictionary), "$dict");
+            bool isInit = false;
+            var cnt = 0;
+            for (var i = 0; i < items.Length; i++) {
+                var item = items[i];
+                if (item.SliceStart is null) {
+                    if (cnt != 0) {
+                        var dict = ReduceDictionaryWithItems(items.Slice(i - cnt, cnt));
+                        if (!isInit) {
+                            expressions.Add(Expression.Assign(varExpr, dict));
+                            isInit = true;
+                        } else {
+                            expressions.Add(Expression.Call(AstMethods.DictUpdate, context, varExpr, dict));
+                        }
+                        cnt = 0;
+                    }
+                    if (!isInit) {
+                        expressions.Add(Expression.Assign(varExpr, EmptyDictExpression));
+                        isInit = true;
+                    }
+                    expressions.Add(Expression.Call(AstMethods.DictUpdate, context, varExpr, TransformOrConstantNull(item.SliceStop, typeof(object))));
+                } else {
+                    cnt++;
+                }
+            }
+            if (cnt != 0) {
+                var dict = ReduceDictionaryWithItems(items.Slice(items.Length - cnt, cnt));
+                if (isInit) {
+                    expressions.Add(Expression.Call(AstMethods.DictUpdate, context, varExpr, dict));
+                } else {
+                    return dict;
+                }
+            }
+            expressions.Add(varExpr);
+            return Expression.Block(typeof(PythonDictionary), new MSAst.ParameterExpression[] { varExpr }, expressions);
         }
 
-        private MSAst.Expression ReduceDictionaryWithItems() {
-            MSAst.Expression[] parts = new MSAst.Expression[_items.Length * 2];
-            Type t = null;
+        private static MSAst.Expression ReduceDictionaryWithItems(ReadOnlySpan<SliceExpression> items) {
+            MSAst.Expression[] parts = new MSAst.Expression[items.Length * 2];
+            Type? t = null;
             bool heterogeneous = false;
-            for (int index = 0; index < _items.Length; index++) {
-                SliceExpression slice = _items[index];
+            for (int index = 0; index < items.Length; index++) {
+                SliceExpression slice = items[index];
                 // Eval order should be:
                 //   { 2 : 1, 4 : 3, 6 :5 }
                 // This is backwards from parameter list eval, so create temporaries to swap ordering.
 
-
                 parts[index * 2] = TransformOrConstantNull(slice.SliceStop, typeof(object));
                 MSAst.Expression key = parts[index * 2 + 1] = TransformOrConstantNull(slice.SliceStart, typeof(object));
 
@@ -68,19 +116,19 @@ private MSAst.Expression ReduceDictionaryWithItems() {
                 }
             }
 
-            return Ast.Call(
+            return Expression.Call(
                 heterogeneous ? AstMethods.MakeDictFromItems : AstMethods.MakeHomogeneousDictFromItems,
-                Ast.NewArrayInit(
+                Expression.NewArrayInit(
                     typeof(object),
                     parts
                 )
             );
         }
 
-        private MSAst.Expression ReduceConstant() {
+        private MSAst.Expression? ReduceConstant() {
             for (int index = 0; index < _items.Length; index++) {
                 SliceExpression slice = _items[index];
-                if (!slice.SliceStop.IsConstant || !slice.SliceStart.IsConstant) {
+                if (slice.SliceStart is null || !slice.SliceStart.IsConstant || !slice.SliceStop!.IsConstant) {
                     return null;
                 }
             }
@@ -89,11 +137,11 @@ private MSAst.Expression ReduceConstant() {
             for (int index = 0; index < _items.Length; index++) {
                 SliceExpression slice = _items[index];
 
-                storage.AddNoLock(slice.SliceStart.GetConstantValue(), slice.SliceStop.GetConstantValue());
+                Debug.Assert(slice.SliceStart is not null);
+                storage.AddNoLock(slice.SliceStart!.GetConstantValue(), slice.SliceStop!.GetConstantValue());
             }
 
-
-            return Ast.Call(AstMethods.MakeConstantDict, Ast.Constant(new ConstantDictionaryStorage(storage), typeof(object)));
+            return Expression.Call(AstMethods.MakeConstantDict, Expression.Constant(new ConstantDictionaryStorage(storage), typeof(object)));
         }
 
         public override void Walk(PythonWalker walker) {
@@ -120,8 +168,8 @@ void IInstructionProvider.AddInstructions(LightCompiler compiler) {
 
         #endregion
 
-        private class EmptyDictInstruction: Instruction {
-            public static EmptyDictInstruction Instance = new EmptyDictInstruction();
+        private class EmptyDictInstruction : Instruction {
+            public static readonly EmptyDictInstruction Instance = new EmptyDictInstruction();
 
             public override int Run(InterpretedFrame frame) {
                 frame.Push(PythonOps.MakeEmptyDict());
diff --git a/Src/IronPython/Compiler/Parser.cs b/Src/IronPython/Compiler/Parser.cs
@@ -2562,8 +2562,10 @@ private IfStatement ParseGenExprIf() {
         }
 
         // dict_display: '{' [dictorsetmaker] '}'
-        // dictorsetmaker: ( (test ':' test (comp_for | (',' test ':' test)* [','])) |
-        //                   (test (comp_for | (',' test)* [','])) )
+        // dictorsetmaker: ( ((test ':' test | '**' expr)
+        //                    (comp_for | (',' (test ':' test | '**' expr))* [','])) |
+        //                   ((test | star_expr)
+        //                    (comp_for | (',' (test | star_expr))* [','])) )
         private Expression FinishDictOrSetValue() {
             var oStart = GetStart();
             var oEnd = GetEnd();
@@ -2578,43 +2580,55 @@ private Expression FinishDictOrSetValue() {
                         break;
                     }
                     bool first = false;
-                    Expression e1 = ParseTest();
-                    if (MaybeEat(TokenKind.Colon)) { // dict literal
-                        if (setMembers != null) {
-                            ReportSyntaxError("invalid syntax");
-                        } else if (dictMembers == null) {
-                            dictMembers = new List<SliceExpression>();
+                    if (MaybeEat(TokenKind.Power)) {
+                        if (setMembers is not null) ReportSyntaxError("invalid syntax");
+                        if (dictMembers is null) {
+                            dictMembers = new();
                             first = true;
                         }
-                        Expression e2 = ParseTest();
-
-                        if (PeekToken(Tokens.KeywordForToken)) {
-                            if (!first) {
+                        var expr = ParseExpr();
+                        var se = new SliceExpression(null, expr, null);
+                        se.SetLoc(_globalParent, expr.StartIndex, expr.EndIndex);
+                        dictMembers.Add(se);
+                    } else {
+                        Expression e1 = ParseTest();
+                        if (MaybeEat(TokenKind.Colon)) { // dict literal
+                            if (setMembers != null) {
                                 ReportSyntaxError("invalid syntax");
+                            } else if (dictMembers == null) {
+                                dictMembers = new List<SliceExpression>();
+                                first = true;
                             }
-                            return FinishDictComp(e1, e2, oStart, oEnd);
-                        }
+                            Expression e2 = ParseTest();
 
-                        SliceExpression se = new SliceExpression(e1, e2, null);
-                        se.SetLoc(_globalParent, e1.StartIndex, e2.EndIndex);
-                        dictMembers.Add(se);
-                    } else { // set literal
-                        if (dictMembers != null) {
-                            ReportSyntaxError("invalid syntax");
-                        } else if (setMembers == null) {
-                            setMembers = new List<Expression>();
-                            first = true;
-                        }
+                            if (PeekToken(Tokens.KeywordForToken)) {
+                                if (!first) {
+                                    ReportSyntaxError("invalid syntax");
+                                }
+                                return FinishDictComp(e1, e2, oStart, oEnd);
+                            }
 
-                        if (PeekToken(Tokens.KeywordForToken)) {
-                            if (!first) {
+                            SliceExpression se = new SliceExpression(e1, e2, null);
+                            se.SetLoc(_globalParent, e1.StartIndex, e2.EndIndex);
+                            dictMembers.Add(se);
+                        } else { // set literal
+                            if (dictMembers != null) {
                                 ReportSyntaxError("invalid syntax");
+                            } else if (setMembers == null) {
+                                setMembers = new List<Expression>();
+                                first = true;
                             }
-                            return FinishSetComp(e1, oStart, oEnd);
-                        }
 
-                        // error recovery
-                        setMembers?.Add(e1);
+                            if (PeekToken(Tokens.KeywordForToken)) {
+                                if (!first) {
+                                    ReportSyntaxError("invalid syntax");
+                                }
+                                return FinishSetComp(e1, oStart, oEnd);
+                            }
+
+                            // error recovery
+                            setMembers?.Add(e1);
+                        }
                     }
 
                     if (!MaybeEat(TokenKind.Comma)) {
diff --git a/Src/IronPython/Runtime/Operations/PythonOps.cs b/Src/IronPython/Runtime/Operations/PythonOps.cs
@@ -926,7 +926,7 @@ internal static bool TryInvokeLengthHint(CodeContext context, object? sequence,
             return PythonCalls.Call(func, allArgs.GetObjectArray());
         }
 
-        public static object GetIndex(CodeContext/*!*/ context, object? o, object index) {
+        public static object GetIndex(CodeContext/*!*/ context, object? o, object? index) {
             PythonContext pc = context.LanguageContext;
             return pc.GetIndexSite.Target(pc.GetIndexSite, o, index);
         }
@@ -1486,6 +1486,26 @@ public static PythonTuple MakeTupleFromSequence(object items) {
             return PythonTuple.Make(items);
         }
 
+        /// <summary>
+        /// DICT_UPDATE
+        /// </summary>
+        [EditorBrowsable(EditorBrowsableState.Never)]
+        public static void DictUpdate(CodeContext context, PythonDictionary dict, object? item) {
+            // call dict.keys()
+            if (!PythonTypeOps.TryInvokeUnaryOperator(context, item, "keys", out object keys)) {
+                throw TypeError($"'{PythonTypeOps.GetName(item)}' object is not a mapping");
+            }
+
+            PythonDictionary res = new PythonDictionary();
+
+            // enumerate the keys getting their values
+            IEnumerator enumerator = GetEnumerator(keys);
+            while (enumerator.MoveNext()) {
+                object? o = enumerator.Current;
+                dict[o] = PythonOps.GetIndex(context, item, o);
+            }
+        }
+
         /// <summary>
         /// Python Runtime Helper for enumerator unpacking (tuple assignments, ...)
         /// Creates enumerator from the input parameter e, and then extracts 
diff --git a/Tests/test_unpack.py b/Tests/test_unpack.py
@@ -194,6 +194,14 @@ def test_ipy3_gh841(self):
         self.assertEqual(a, 1)
         self.assertEqual(b, [2])
 
+    def test_unpack_dict(self):
+        self.assertEqual({'x': 1, **{'y': 2}}, {'x': 1, 'y': 2})
+        self.assertEqual({'x': 1, **{'x': 2}}, {'x': 2})
+        self.assertEqual({**{'x': 2}, 'x': 1}, {'x': 1})
+
+        with self.assertRaises(TypeError):
+            {**[]}
+
 def test_main():
     test.support.run_unittest(UnpackTest)
 
