Add the ability to change Vi cmd mode dispatch tables.

Author: jazzdelightsme

PSReadLine/Cmdlets.cs

@@ -52,6 +52,12 @@ public enum ViModeStyle
         Prompt,
         Cursor
     }
+
+    public enum ViMode
+    {
+        Insert,
+        Command
+    }
     #endregion vi
 
     public enum HistorySaveStyle
@@ -590,15 +596,48 @@ protected override void EndProcessing()
         }
     }
 
-    [Cmdlet("Set", "PSReadlineKeyHandler", HelpUri = "http://go.microsoft.com/fwlink/?LinkId=528810")]
-    public class SetPSReadlineKeyHandlerCommand : PSCmdlet, IDynamicParameters
+    public class ChangePSReadlineKeyHandlerCommand : PSCmdlet
     {
         [Parameter(Position = 0, Mandatory = true)]
         [Alias("Key")]
         [ValidateNotNullOrEmpty]
         [SuppressMessage("Microsoft.Performance", "CA1819:PropertiesShouldNotReturnArrays")]
         public string[] Chord { get; set; }
 
+        [Parameter]
+        public ViMode ViMode { get; set; }
+
+        [ExcludeFromCodeCoverage]
+        protected IDisposable UseRequestedDispatchTables()
+        {
+            bool inViMode = PSConsoleReadLine.GetOptions().EditMode == EditMode.Vi;
+            bool viModeParamPresent = MyInvocation.BoundParameters.ContainsKey("ViMode");
+
+            if (inViMode || viModeParamPresent)
+            {
+                if (!inViMode)
+                {
+                    // "-ViMode" must have been specified explicitly. Well, okay... we can
+                    // modify the Vi tables... but isn't that an odd thing to do from
+                    // not-vi mode?
+                    WriteWarning(PSReadLineResources.NotInViMode);
+                }
+
+                if (ViMode == ViMode.Insert) // default if -ViMode not specified
+                    return PSConsoleReadLine.UseViInsertModeTables();
+                else if (ViMode == ViMode.Command)
+                    return PSConsoleReadLine.UseViCommandModeTables();
+                else
+                    System.Diagnostics.Debug.Fail("unexpected ViMode");
+            }
+
+            return null;
+        }
+    }
+
+    [Cmdlet("Set", "PSReadlineKeyHandler", HelpUri = "http://go.microsoft.com/fwlink/?LinkId=528810")]
+    public class SetPSReadlineKeyHandlerCommand : ChangePSReadlineKeyHandlerCommand, IDynamicParameters
+    {
         [Parameter(Position = 1, Mandatory = true, ParameterSetName = "ScriptBlock")]
         [ValidateNotNull]
         public ScriptBlock ScriptBlock { get; set; }
@@ -616,18 +655,21 @@ public class SetPSReadlineKeyHandlerCommand : PSCmdlet, IDynamicParameters
         [ExcludeFromCodeCoverage]
         protected override void EndProcessing()
         {
-            if (ParameterSetName.Equals(FunctionParameterSet))
-            {
-                var function = (string)_dynamicParameters.Value[FunctionParameter].Value;
-                var keyHandler = (Action<ConsoleKeyInfo?, object>)
-                    Delegate.CreateDelegate(typeof (Action<ConsoleKeyInfo?, object>),
-                        typeof (PSConsoleReadLine).GetMethod(function));
-                BriefDescription = function;
-                PSConsoleReadLine.SetKeyHandler(Chord, keyHandler, BriefDescription, Description);
-            }
-            else
+            using (UseRequestedDispatchTables())
             {
-                PSConsoleReadLine.SetKeyHandler(Chord, ScriptBlock, BriefDescription, Description);
+                if (ParameterSetName.Equals(FunctionParameterSet))
+                {
+                    var function = (string)_dynamicParameters.Value[FunctionParameter].Value;
+                    var keyHandler = (Action<ConsoleKeyInfo?, object>)
+                        Delegate.CreateDelegate(typeof (Action<ConsoleKeyInfo?, object>),
+                            typeof (PSConsoleReadLine).GetMethod(function));
+                    BriefDescription = function;
+                    PSConsoleReadLine.SetKeyHandler(Chord, keyHandler, BriefDescription, Description);
+                }
+                else
+                {
+                    PSConsoleReadLine.SetKeyHandler(Chord, ScriptBlock, BriefDescription, Description);
+                }
             }
         }
 
@@ -713,18 +755,15 @@ protected override void EndProcessing()
     }
 
     [Cmdlet("Remove", "PSReadlineKeyHandler", HelpUri = "http://go.microsoft.com/fwlink/?LinkId=528809")]
-    public class RemoveKeyHandlerCommand : PSCmdlet
+    public class RemoveKeyHandlerCommand : ChangePSReadlineKeyHandlerCommand
     {
-        [Parameter(Position = 0, Mandatory = true)]
-        [Alias("Key")]
-        [ValidateNotNullOrEmpty]
-        [SuppressMessage("Microsoft.Performance", "CA1819:PropertiesShouldNotReturnArrays")]
-        public string[] Chord { get; set; }
-
         [ExcludeFromCodeCoverage]
         protected override void EndProcessing()
         {
-            PSConsoleReadLine.RemoveKeyHandler(Chord);
+            using (UseRequestedDispatchTables())
+            {
+                PSConsoleReadLine.RemoveKeyHandler(Chord);
+            }
         }
     }
 

PSReadLine/Disposable.cs

@@ -0,0 +1,27 @@
+using System;
+
+namespace Microsoft.PowerShell
+{
+    internal sealed class Disposable : IDisposable
+    {
+        private Action m_onDispose;
+
+        public Disposable(Action onDispose)
+        {
+            if (onDispose == null)
+                throw new ArgumentNullException("onDispose");
+
+            m_onDispose = onDispose;
+        }
+
+        public void Dispose()
+        {
+            if (m_onDispose != null)
+            {
+                m_onDispose();
+                m_onDispose = null;
+            }
+        }
+    }
+}
+

PSReadLine/PSReadLine.csproj

@@ -57,6 +57,7 @@
     <Compile Include="ConsoleBufferBuilder.cs" />
     <Compile Include="ConsoleKeyChordConverter.cs" />
     <Compile Include="ConsoleLib.cs" />
+    <Compile Include="Disposable.cs" />
     <Compile Include="History.cs" />
     <Compile Include="HistoryQueue.cs" />
     <Compile Include="KeyBindings.cs" />

PSReadLine/PSReadLineResources.Designer.cs

@@ -717,4 +717,7 @@ If there are other parse errors, unresolved commands, or incorrect parameters, s
   <data name="ViJoinLinesDescription" xml:space="preserve">
     <value>Joins the current multi-line edit mode line with the next.</value>
   </data>
+  <data name="NotInViMode" xml:space="preserve">
+    <value>The -ViMode parameter was used, but the current EditMode is not Vi.</value>
+  </data>
 </root>

PSReadLine/PSReadLineResources.resx

@@ -511,6 +511,42 @@ public static void ViInsertMode(ConsoleKeyInfo? key = null, object arg = null)
             _singleton.ViIndicateInsertMode();
         }
 
+        /// <summary>
+        /// Temporarily swap in Vi-Command dispatch tables. Used for setting handlers.
+        /// </summary>
+        internal static IDisposable UseViCommandModeTables()
+        {
+            var oldDispatchTable = _singleton._dispatchTable;
+            var oldChordDispatchTable = _singleton._chordDispatchTable;
+
+            _singleton._dispatchTable = _viCmdKeyMap;
+            _singleton._chordDispatchTable = _viCmdChordTable;
+
+            return new Disposable( () =>
+            {
+                _singleton._dispatchTable = oldDispatchTable;
+                _singleton._chordDispatchTable = oldChordDispatchTable;
+            } );
+        }
+
+        /// <summary>
+        /// Temporarily swap in Vi-Insert dispatch tables. Used for setting handlers.
+        /// </summary>
+        internal static IDisposable UseViInsertModeTables()
+        {
+            var oldDispatchTable = _singleton._dispatchTable;
+            var oldChordDispatchTable = _singleton._chordDispatchTable;
+
+            _singleton._dispatchTable = _viInsKeyMap;
+            _singleton._chordDispatchTable = _viInsChordTable;
+
+            return new Disposable( () =>
+            {
+                _singleton._dispatchTable = oldDispatchTable;
+                _singleton._chordDispatchTable = oldChordDispatchTable;
+            } );
+        }
+
         private void ViIndicateCommandMode()
         {
             if (_options.ViModeIndicator == ViModeStyle.Cursor)