New jupyter notebook fixes

src/main/java/com/maddyhome/idea/vim/EventFacade.java

@@ -27,6 +27,7 @@ import com.intellij.openapi.editor.EditorFactory;
 import com.intellij.openapi.editor.actionSystem.TypedAction;
 import com.intellij.openapi.editor.actionSystem.TypedActionHandler;
 import com.intellij.openapi.editor.event.*;
+import com.maddyhome.idea.vim.helper.HandlerInjector;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -51,6 +52,15 @@ public class EventFacade {
 
   public void setupTypedActionHandler(@NotNull VimTypedActionHandler handler) {
     final TypedAction typedAction = getTypedAction();
+
+    if (HandlerInjector.notebookCommandMode()) {
+      TypedActionHandler result = HandlerInjector.inject();
+      if (result != null) {
+        myOriginalTypedActionHandler = result;
+        return;
+      }
+    }
+
     myOriginalTypedActionHandler = typedAction.getRawHandler();
 
     typedAction.setupRawHandler(handler);

src/main/java/com/maddyhome/idea/vim/action/VimShortcutKeyAction.kt

@@ -24,6 +24,7 @@ import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.EmptyAction
 import com.intellij.openapi.actionSystem.PlatformDataKeys
+import com.intellij.openapi.application.invokeLater
 import com.intellij.openapi.diagnostic.debug
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.editor.Editor
@@ -35,12 +36,14 @@ import com.maddyhome.idea.vim.KeyHandler
 import com.maddyhome.idea.vim.VimPlugin
 import com.maddyhome.idea.vim.helper.EditorDataContext
 import com.maddyhome.idea.vim.helper.EditorHelper
+import com.maddyhome.idea.vim.helper.HandlerInjector
 import com.maddyhome.idea.vim.helper.StringHelper
 import com.maddyhome.idea.vim.helper.inInsertMode
 import com.maddyhome.idea.vim.helper.inNormalMode
 import com.maddyhome.idea.vim.helper.isIdeaVimDisabledHere
 import com.maddyhome.idea.vim.helper.isPrimaryEditor
 import com.maddyhome.idea.vim.helper.isTemplateActive
+import com.maddyhome.idea.vim.helper.updateCaretsVisualAttributes
 import com.maddyhome.idea.vim.key.ShortcutOwner
 import com.maddyhome.idea.vim.key.ShortcutOwnerInfo
 import com.maddyhome.idea.vim.listener.AceJumpService
@@ -116,13 +119,21 @@ class VimShortcutKeyAction : AnAction(), DumbAware/*, LightEditCompatible*/ {
         return false
       }
 
+      val keyCode = keyStroke.keyCode
+
+      if (HandlerInjector.notebookCommandMode()) {
+        LOG.trace("Python Notebook command mode")
+        if (keyCode == KeyEvent.VK_RIGHT || keyCode == KeyEvent.VK_KP_RIGHT || keyCode == KeyEvent.VK_ENTER) {
+          invokeLater { editor.updateCaretsVisualAttributes() }
+        }
+        return false
+      }
+
       if (AceJumpService.getInstance()?.isActive(editor) == true) {
         LOG.trace("Do not execute shortcut because AceJump is active")
         return false
       }
 
-      val keyCode = keyStroke.keyCode
-
       if (LookupManager.getActiveLookup(editor) != null && !LookupKeys.isEnabledForLookup(keyStroke)) {
         LOG.trace("Do not execute shortcut because of lookup keys")
         return false

src/main/java/com/maddyhome/idea/vim/helper/HandlerInjector.kt

@@ -0,0 +1,58 @@
+/*
+ * IdeaVim - Vim emulator for IDEs based on the IntelliJ platform
+ * Copyright (C) 2003-2021 The IdeaVim authors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.maddyhome.idea.vim.helper
+
+import com.intellij.openapi.editor.actionSystem.TypedAction
+import com.intellij.openapi.editor.actionSystem.TypedActionHandler
+import com.maddyhome.idea.vim.VimTypedActionHandler
+import java.lang.reflect.Field
+import java.lang.reflect.Modifier
+
+class HandlerInjector {
+  companion object {
+    @JvmStatic
+    fun inject(): TypedActionHandler? {
+      try {
+        val javaClass = TypedAction.getInstance().rawHandler::class.java
+        val pythonHandler = javaClass.kotlin.objectInstance
+        val field =
+          javaClass.declaredFields.singleOrNull { it.name == "DEFAULT_RAW_HANDLER" || it.name == "editModeRawHandler" }
+            ?: return null
+        field.isAccessible = true
+
+        val modifiersField: Field = Field::class.java.getDeclaredField("modifiers")
+        modifiersField.isAccessible = true
+        modifiersField.setInt(field, field.modifiers and Modifier.FINAL.inv())
+
+        val originalHandler = field.get(pythonHandler) as? TypedActionHandler ?: return null
+        val newVimHandler = VimTypedActionHandler(originalHandler)
+        field.set(pythonHandler, newVimHandler)
+        return originalHandler
+      } catch (ignored: Exception) {
+        // Ignore
+      }
+      return null
+    }
+
+    @JvmStatic
+    fun notebookCommandMode(): Boolean {
+      return TypedAction.getInstance().rawHandler::class.java.simpleName.equals("JupyterCommandModeTypingBlocker")
+    }
+  }
+}