Remember script block for custom key bindings

For various reasons, we may need to know that a key binding
came from a script block so that is now saved in the key binding.

Author: lzybkr

PSReadLine/Cmdlets.cs

@@ -513,20 +513,19 @@ public class SetPSReadlineKeyHandlerCommand : PSCmdlet, IDynamicParameters
         [ExcludeFromCodeCoverage]
         protected override void EndProcessing()
         {
-            Action<ConsoleKeyInfo?, object> keyHandler;
             if (ParameterSetName.Equals(FunctionParameterSet))
             {
                 var function = (string)_dynamicParameters.Value[FunctionParameter].Value;
-                keyHandler = (Action<ConsoleKeyInfo?, object>)
+                var keyHandler = (Action<ConsoleKeyInfo?, object>)
                     Delegate.CreateDelegate(typeof (Action<ConsoleKeyInfo?, object>),
-                                            typeof (PSConsoleReadLine).GetMethod(function));
+                        typeof (PSConsoleReadLine).GetMethod(function));
                 BriefDescription = function;
+                PSConsoleReadLine.SetKeyHandler(Chord, keyHandler, BriefDescription, Description);
             }
             else
             {
-                keyHandler = (key, arg) => ScriptBlock.Invoke(key, arg);
+                PSConsoleReadLine.SetKeyHandler(Chord, ScriptBlock, BriefDescription, Description);
             }
-            PSConsoleReadLine.SetKeyHandler(Chord, keyHandler, BriefDescription, Description);
         }
 
         private readonly Lazy<RuntimeDefinedParameterDictionary> _dynamicParameters =

PSReadLine/KeyBindings.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Management.Automation;
 using System.Text;
 
 namespace PSConsoleUtilities
@@ -43,6 +44,7 @@ class KeyHandler
             public Action<ConsoleKeyInfo?, object> Action;
             public string BriefDescription;
             public string LongDescription;
+            public ScriptBlock ScriptBlock;
         }
 
         internal class ConsoleKeyInfoComparer : IEqualityComparer<ConsoleKeyInfo>
@@ -58,7 +60,7 @@ public int GetHashCode(ConsoleKeyInfo obj)
             }
         }
 
-        static KeyHandler MakeKeyHandler(Action<ConsoleKeyInfo?, object> action, string briefDescription, string longDescription = null)
+        static KeyHandler MakeKeyHandler(Action<ConsoleKeyInfo?, object> action, string briefDescription, string longDescription = null, ScriptBlock scriptBlock = null)
         {
             if (string.IsNullOrWhiteSpace(longDescription))
                 longDescription = PSReadLineResources.ResourceManager.GetString(briefDescription + "Description");
@@ -67,7 +69,8 @@ static KeyHandler MakeKeyHandler(Action<ConsoleKeyInfo?, object> action, string
             {
                 Action = action,
                 BriefDescription = briefDescription,
-                LongDescription = longDescription
+                LongDescription = longDescription,
+                ScriptBlock = scriptBlock,
             };
         }
 

PSReadLine/Options.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Management.Automation;
 using System.Reflection;
 using System.Threading;
 
@@ -155,7 +156,7 @@ private void SetOptionsInternal(SetPSReadlineOption options)
             }
         }
 
-        private void SetKeyHandlerInternal(string[] keys, Action<ConsoleKeyInfo?, object> handler, string briefDescription, string longDescription)
+        private void SetKeyHandlerInternal(string[] keys, Action<ConsoleKeyInfo?, object> handler, string briefDescription, string longDescription, ScriptBlock scriptBlock)
         {
             foreach (var key in keys)
             {
@@ -196,12 +197,22 @@ public static PSConsoleReadlineOptions GetOptions()
             return _singleton._options;
         }
 
+        /// <summary>
+        /// Helper function for the Set-PSReadlineKeyHandler cmdlet.
+        /// </summary>
+        public static void SetKeyHandler(string[] key, ScriptBlock scriptBlock, string briefDescription, string longDescription)
+        {
+            Action<ConsoleKeyInfo?, object> handler =
+                (k, arg) => scriptBlock.Invoke(k, arg);
+            _singleton.SetKeyHandlerInternal(key, handler, briefDescription, longDescription, scriptBlock);
+        }
+
         /// <summary>
         /// Helper function for the Set-PSReadlineKeyHandler cmdlet.
         /// </summary>
         public static void SetKeyHandler(string[] key, Action<ConsoleKeyInfo?, object> handler, string briefDescription, string longDescription)
         {
-            _singleton.SetKeyHandlerInternal(key, handler, briefDescription, longDescription);
+            _singleton.SetKeyHandlerInternal(key, handler, briefDescription, longDescription, null);
         }
 
         /// <summary>