PY-77951 [Jupyter] Fix cell detection for action data

(cherry picked from commit b420ce218b422830770bd6f7fc0021628b4acee8)

GitOrigin-RevId: a43a5f964d41f459f801809047f33eac6cf5c35d

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/DecoratedEditor.kt

@@ -3,9 +3,12 @@ package com.intellij.notebooks.visualization.ui
 import com.intellij.notebooks.ui.editor.actions.command.mode.NotebookEditorMode
 import com.intellij.notebooks.ui.editor.actions.command.mode.setMode
 import com.intellij.notebooks.visualization.*
+import com.intellij.notebooks.visualization.context.NotebookDataContext.NOTEBOOK_CELL_LINES_INTERVAL
 import com.intellij.notebooks.visualization.inlay.JupyterBoundsChangeHandler
 import com.intellij.notebooks.visualization.ui.EditorCellViewEventListener.CellViewRemoved
 import com.intellij.notebooks.visualization.ui.EditorCellViewEventListener.EditorCellViewEvent
+import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.WriteIntentReadAction
 import com.intellij.openapi.client.ClientSystemInfo
@@ -23,6 +26,7 @@ import java.awt.GraphicsEnvironment
 import java.awt.Point
 import java.awt.event.InputEvent
 import java.awt.event.MouseEvent
+import java.awt.event.MouseEvent.MOUSE_PRESSED
 import java.awt.event.MouseWheelEvent
 import java.util.concurrent.atomic.AtomicBoolean
 import javax.swing.*
@@ -35,7 +39,8 @@ class DecoratedEditor private constructor(
 ) : NotebookEditor {
 
   /** Used to hold current cell under mouse, to update the folding state and "run" button state. */
-  private var mouseOverCell: EditorCellView? = null
+  override var mouseOverCell: EditorCellView? = null
+    private set
 
   private val selectionModel = EditorCellSelectionModel(manager)
 
@@ -58,20 +63,6 @@ class DecoratedEditor private constructor(
       }
     }, editorImpl.disposable)
 
-    editorImpl.addEditorMouseListener(object : EditorMouseListener {
-      override fun mousePressed(event: EditorMouseEvent) {
-        if (!event.isConsumed && event.mouseEvent.button == MouseEvent.BUTTON1) {
-          val point = getEditorPoint(event.mouseEvent)?.second ?: return
-
-          val selectedCell = getCellViewByPoint(point)?.cell ?: return
-
-          if (event.area != EditorMouseEventArea.EDITING_AREA) {
-            mousePressed(selectedCell.interval, event.isCtrlPressed(), event.isShiftPressed())
-          }
-        }
-      }
-    }, editorImpl.disposable)
-
     editorImpl.caretModel.addCaretListener(object : CaretListener {
       override fun caretAdded(event: CaretEvent) = scheduleSelectionUpdate()
       override fun caretPositionChanged(event: CaretEvent) = scheduleSelectionUpdate()
@@ -101,7 +92,7 @@ class DecoratedEditor private constructor(
           nestedScrollingSupport.processMouseWheelEvent(event)
         }
         else if (event is MouseEvent) {
-          if (event.id == MouseEvent.MOUSE_CLICKED || event.id == MouseEvent.MOUSE_RELEASED || event.id == MouseEvent.MOUSE_PRESSED) {
+          if (event.id == MouseEvent.MOUSE_CLICKED || event.id == MouseEvent.MOUSE_RELEASED || event.id == MOUSE_PRESSED) {
             ComponentUtil.getParentOfType(JScrollPane::class.java, (event.component as? JComponent)
               ?.findComponentAt(event.point))
               ?.let { scrollPane ->
@@ -114,6 +105,19 @@ class DecoratedEditor private constructor(
         }
       }
 
+      eventDispatcher.addListener { event ->
+        if (event.id == MOUSE_PRESSED && event is MouseEvent) {
+          val point = getEditorPoint(event)?.second ?: return@addListener
+
+          val selectedCell = getCellViewByPoint(point)?.cell ?: return@addListener
+
+          if (editorImpl.getMouseEventArea(event) != EditorMouseEventArea.EDITING_AREA) {
+            editorImpl.setMode(NotebookEditorMode.COMMAND)
+          }
+          updateSelectionAfterClick(selectedCell.interval, event.isCtrlPressed(), event.isShiftPressed(), event.button)
+        }
+      }
+
       Disposer.register(editor.disposable, this)
     }
 
@@ -121,7 +125,11 @@ class DecoratedEditor private constructor(
   }
 
   /** The main thing while we need it - to perform updating of underlying components within keepScrollingPositionWhile. */
-  private class EditorComponentWrapper(private val editor: Editor, private val editorViewport: JViewport, component: Component) : JPanel(BorderLayout()) {
+  private class EditorComponentWrapper(
+    private val editor: Editor,
+    private val editorViewport: JViewport,
+    component: Component,
+  ) : JPanel(BorderLayout()), UiDataProvider {
     init {
       isOpaque = false
       // The reason why we need to wrap into fate viewport is the code in [com/intellij/openapi/editor/impl/EditorImpl.java:2031]
@@ -141,6 +149,10 @@ class DecoratedEditor private constructor(
         JupyterBoundsChangeHandler.get(editor).performPostponed()
       }
     }
+
+    override fun uiDataSnapshot(sink: DataSink) {
+      sink[NOTEBOOK_CELL_LINES_INTERVAL] = editor.notebookEditor.mouseOverCell?.cell?.interval
+    }
   }
 
   private fun scheduleSelectionUpdate() {
@@ -214,17 +226,13 @@ class DecoratedEditor private constructor(
     return cur?.view
   }
 
-  override fun inlayClicked(clickedCell: NotebookCellLines.Interval, ctrlPressed: Boolean, shiftPressed: Boolean) {
-    mousePressed(clickedCell, ctrlPressed, shiftPressed)
-  }
-
-  private fun mousePressed(clickedCell: NotebookCellLines.Interval, ctrlPressed: Boolean, shiftPressed: Boolean) {
+  override fun inlayClicked(clickedCell: NotebookCellLines.Interval, ctrlPressed: Boolean, shiftPressed: Boolean, mouseButton: Int) {
     editorImpl.setMode(NotebookEditorMode.COMMAND)
-    updateSelectionAfterClick(clickedCell, ctrlPressed, shiftPressed)
+    updateSelectionAfterClick(clickedCell, ctrlPressed, shiftPressed, mouseButton)
   }
 
   @Suppress("ConvertArgumentToSet")
-  private fun updateSelectionAfterClick(clickedCell: NotebookCellLines.Interval, ctrlPressed: Boolean, shiftPressed: Boolean) {
+  private fun updateSelectionAfterClick(clickedCell: NotebookCellLines.Interval, ctrlPressed: Boolean, shiftPressed: Boolean, mouseButton: Int) {
     val model = editorImpl.cellSelectionModel!!
     when {
       ctrlPressed -> {
@@ -260,7 +268,9 @@ class DecoratedEditor private constructor(
         }
       }
       else -> {
-        model.selectSingleCell(clickedCell)
+        if (mouseButton == MouseEvent.BUTTON1 || !model.isSelectedCell(clickedCell)) {
+          model.selectSingleCell(clickedCell)
+        }
       }
     }
   }
@@ -290,8 +300,8 @@ internal fun <T> keepScrollingPositionWhile(editor: Editor, task: () -> T): T {
 private fun hasIntersection(cells: List<NotebookCellLines.Interval>, others: List<NotebookCellLines.Interval>): Boolean =
   !(cells.last().ordinal < others.first().ordinal || cells.first().ordinal > others.last().ordinal)
 
-private fun EditorMouseEvent.isCtrlPressed(): Boolean =
-  (mouseEvent.modifiersEx and if (ClientSystemInfo.isMac()) InputEvent.META_DOWN_MASK else InputEvent.CTRL_DOWN_MASK) != 0
+private fun MouseEvent.isCtrlPressed(): Boolean =
+  (modifiersEx and if (ClientSystemInfo.isMac()) InputEvent.META_DOWN_MASK else InputEvent.CTRL_DOWN_MASK) != 0
 
-private fun EditorMouseEvent.isShiftPressed(): Boolean =
-  (mouseEvent.modifiersEx and InputEvent.SHIFT_DOWN_MASK) != 0
+private fun MouseEvent.isShiftPressed(): Boolean =
+  (modifiersEx and InputEvent.SHIFT_DOWN_MASK) != 0

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/NotebookEditor.kt

@@ -5,7 +5,8 @@ import com.intellij.openapi.util.Key
 import com.intellij.notebooks.visualization.NotebookCellLines
 
 interface NotebookEditor {
-  fun inlayClicked(clickedCell: NotebookCellLines.Interval, ctrlPressed: Boolean, shiftPressed: Boolean)
+  val mouseOverCell: EditorCellView?
+  fun inlayClicked(clickedCell: NotebookCellLines.Interval, ctrlPressed: Boolean, shiftPressed: Boolean, mouseButton: Int)
 }
 
 internal val notebookEditorKey = Key.create<NotebookEditor>(NotebookEditor::class.java.name)