fix bug with remove/removeKeyFromDictionary

ServiceStack · Apr 11, 2020 · 3975f86 · 3975f86
1 parent 4ee0482
commit 3975f86
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/src/ServiceStack.Common/Script/Methods/DefaultScripts.cs b/src/ServiceStack.Common/Script/Methods/DefaultScripts.cs
@@ -1313,11 +1313,11 @@ public async Task selectPartial(ScriptScopeContext scope, object target, string
 
         public object removeKeyFromDictionary(IDictionary dictionary, object keyToRemove)
         {
-            var removeKeys = keyToRemove is IEnumerable e
+            var removeKeys = keyToRemove is IEnumerable e && !(keyToRemove is string)
                 ? e.Map(x => x)
                 : null;
 
-            foreach (var key in dictionary.Keys)
+            foreach (var key in EnumerableUtils.ToList(dictionary.Keys))
             {
                 if (removeKeys != null)
                 {
@@ -1337,11 +1337,12 @@ public object removeKeyFromDictionary(IDictionary dictionary, object keyToRemove
 
         public object remove(object target, object keysToRemove)
         {
-            var removeKeys = keysToRemove is IEnumerable eKeys
-                ? eKeys.Map(x => x)
-                : null;
+            var removeKeys = keysToRemove is string s
+                ? (IEnumerable) new[] {s}
+                : keysToRemove is IEnumerable eKeys
+                    ? eKeys.Map(x => x)
+                    : null;
 
-            var stringKey = keysToRemove as string;
             var stringKeys = removeKeys?.OfType<string>().ToArray();
             if (stringKeys.IsEmpty())
                 stringKeys = null;

diff --git a/tests/ServiceStack.WebHost.Endpoints.Tests/ScriptTests/DefaultScriptsTests.cs b/tests/ServiceStack.WebHost.Endpoints.Tests/ScriptTests/DefaultScriptsTests.cs
@@ -2086,5 +2086,21 @@ public void Can_flatMap()
             Assert.That(context.Evaluate("{{ flatten([[1,2],[3,4]]) |> return }}"), Is.EqualTo(new[]{ 1, 2, 3, 4 }));
         }
 
+        [Test]
+        public void Can_removeKeyFromDictionary()
+        {
+            var context = new ScriptContext().Init();
+            var output = context.RenderScript(@"```code|q
+sample = {}
+sample.myKey1 = 1
+sample.myKey2 = 2
+sample |> remove('myKey1')
+sample |> removeKeyFromDictionary('myKey2')
+```
+{{ sample.myKey }}".NormalizeNewLines());
+
+            Assert.That(output, Is.EqualTo(""));
+        }
+
     }
 }