PY-72902 Jupyter(refactor): refactor context, create JupyterDataContext and NotebookDataContext

GitOrigin-RevId: 1794b8f2db92ef6a2b27adb1ccf36cfc09054887

notebooks/notebook-ui/src/com/intellij/notebooks/ui/jupyterToolbar/JupyterAddCellToolbarService.kt

@@ -1,13 +1,13 @@
 package com.intellij.notebooks.ui.jupyterToolbar
 
 import com.intellij.ide.ui.customization.CustomActionsSchema
+import com.intellij.notebooks.ui.visualization.DefaultNotebookEditorAppearanceSizes
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.actionSystem.ActionGroup
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import kotlinx.coroutines.*
-import com.intellij.notebooks.ui.visualization.DefaultNotebookEditorAppearanceSizes
 import java.awt.MouseInfo
 import java.awt.Point
 import java.awt.Rectangle
@@ -97,7 +97,8 @@ class JupyterAddCellToolbarService(private val scope: CoroutineScope): Disposabl
 
   private fun createAndShowToolbar(editor: Editor) {
     actionGroup ?: return
-    if (currentToolbar == null) currentToolbar = JupyterAddNewCellToolbar(actionGroup, editor.contentComponent)
+    if (currentToolbar == null)
+      currentToolbar = JupyterAddNewCellToolbar(actionGroup, editor.contentComponent)
     editor.contentComponent.add(currentToolbar, 0)
     hideToolbarTimer.stop()
     adjustToolbarPosition()

notebooks/notebook-ui/src/com/intellij/notebooks/ui/jupyterToolbar/JupyterAddNewCellToolbar.kt

@@ -4,7 +4,6 @@ import com.intellij.openapi.actionSystem.ActionGroup
 import com.intellij.openapi.actionSystem.ActionPlaces
 import com.intellij.openapi.actionSystem.impl.ActionToolbarImpl
 import com.intellij.openapi.actionSystem.toolbarLayout.ToolbarLayoutStrategy
-import com.intellij.openapi.editor.Editor
 import com.intellij.ui.JBColor
 import com.intellij.ui.NewUiValue
 import com.intellij.ui.RoundedLineBorder
@@ -17,11 +16,11 @@ import java.awt.Graphics
 import java.awt.Graphics2D
 import javax.swing.BorderFactory
 import javax.swing.JComponent
-import javax.swing.SwingUtilities
 
 @ApiStatus.Internal
-class JupyterAddNewCellToolbar(  // PY-66455
-  actionGroup: ActionGroup, target: JComponent, place: String = ActionPlaces.EDITOR_INLAY
+class JupyterAddNewCellToolbar(
+  // PY-66455
+  actionGroup: ActionGroup, target: JComponent, place: String = ActionPlaces.EDITOR_INLAY,
 ) : ActionToolbarImpl(place, actionGroup, true) {
   init {
     val borderColor = when (NewUiValue.isEnabled()) {
@@ -36,6 +35,7 @@ class JupyterAddNewCellToolbar(  // PY-66455
     setSkipWindowAdjustments(false)
   }
 
+
   override fun paintComponent(g: Graphics) {
     val g2 = g.create() as Graphics2D
     try {
@@ -56,14 +56,6 @@ class JupyterAddNewCellToolbar(  // PY-66455
     layoutStrategy = ToolbarLayoutStrategy.NOWRAP_STRATEGY
   }
 
-  fun getRespectiveLineNumberInEditor(editor: Editor): Int? {
-    val point = SwingUtilities.convertPoint(this, 0, this.height, editor.contentComponent)
-    val documentLineCount = editor.document.lineCount
-    var prospectiveLineNumber = point.y.takeIf { it >= 0 }?.let { editor.xyToLogicalPosition(point).line } ?: return null
-    if (prospectiveLineNumber >= documentLineCount) { prospectiveLineNumber = documentLineCount - 1 }
-    return prospectiveLineNumber
-  }
-
   override fun installPopupHandler(customizable: Boolean, popupActionGroup: ActionGroup?, popupActionId: String?) = Unit
 
   companion object {

notebooks/visualization/resources/intellij.notebooks.visualization.xml

@@ -45,10 +45,6 @@
   <extensions defaultExtensionNs="com.intellij">
     <notificationGroup displayType="BALLOON" id="Notebook Table" bundle="messages.VisualizationBundle" key="inlay.output.table.notification.group.name"/>
 
-    <getDataRule
-      key="NOTEBOOK_CELL_LINES_INTERVAL"
-      implementationClass="com.intellij.notebooks.visualization.NotebookCellLinesIntervalDataRule"/>
-    <uiDataRule implementation="com.intellij.notebooks.visualization.EditorsWithOffsetsDataRule"/>
     <editorFactoryDocumentListener implementation="com.intellij.notebooks.visualization.UndoableActionListener"
                                    order="last"/>
   </extensions>

notebooks/visualization/src/com/intellij/notebooks/visualization/CellLinesActionBase.kt

@@ -1,19 +1,18 @@
 package com.intellij.notebooks.visualization
 
+import com.intellij.notebooks.visualization.context.NotebookDataContext.notebookEditor
 import com.intellij.openapi.actionSystem.ActionUpdateThread
 import com.intellij.openapi.actionSystem.AnActionEvent
-import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.actionSystem.PlatformCoreDataKeys
 import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.fileEditor.impl.NonProjectFileWritingAccessProvider
 import com.intellij.openapi.project.DumbAwareAction
-import com.intellij.util.asSafely
 
 abstract class CellLinesActionBase : DumbAwareAction() {
   override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.BGT
 
   override fun update(event: AnActionEvent) {
-    event.presentation.isEnabled = event.getAnyEditor()
+    event.presentation.isEnabled = event.dataContext.notebookEditor
       ?.takeIf { NotebookCellLines.hasSupport(it) } != null
   }
 
@@ -25,19 +24,12 @@ abstract class CellLinesActionBase : DumbAwareAction() {
 
   final override fun actionPerformed(event: AnActionEvent) {
     if (isModifyingSourceCode() && !isEditingAllowed(event)) return
-    val editor = event.getAnyEditor() ?: return
+    val editor = event.dataContext.notebookEditor ?: return
     val cellLines = NotebookCellLines.get(editor)
     actionPerformed(event, editor, cellLines)
   }
 }
 
-fun DataContext.getAnyEditor() =
-  getData(EDITORS_WITH_OFFSETS_DATA_KEY)
-    ?.firstOrNull()
-    ?.first
-    ?.asSafely<EditorImpl>()
-
-fun AnActionEvent.getAnyEditor(): EditorImpl? = dataContext.getAnyEditor()
 
 fun isEditingAllowed(e: AnActionEvent): Boolean {
   val virtualFile = e.dataContext.getData(PlatformCoreDataKeys.FILE_EDITOR)?.file ?: return false

notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookCellLines.kt

@@ -1,16 +1,14 @@
 package com.intellij.notebooks.visualization
 
 import com.intellij.lang.Language
-import com.intellij.openapi.actionSystem.DataKey
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.Key
+import com.intellij.openapi.util.TextRange
 import com.intellij.util.EventDispatcher
 import com.intellij.util.keyFMap.KeyFMap
 import java.util.*
 
-val NOTEBOOK_CELL_LINES_INTERVAL_DATA_KEY = DataKey.create<NotebookCellLines.Interval>("NOTEBOOK_CELL_LINES_INTERVAL")
-
 /**
  * Incrementally iterates over Notebook document, calculates line ranges of cells using lexer.
  * Fast enough for running in EDT, but could be used in any other thread.
@@ -42,9 +40,54 @@ interface NotebookCellLines {
   ) : Comparable<Interval> {
     val language: Language = data.get(INTERVAL_LANGUAGE_KEY)!!
 
+    val firstContentLine: Int
+      get() =
+        if (markers.hasTopLine) lines.first + 1
+        else lines.first
+
+    val lastContentLine: Int
+      get() =
+        if (markers.hasBottomLine) lines.last - 1
+        else lines.last
+
+    val contentLines: IntRange
+      get() = firstContentLine..lastContentLine
+
+
     operator fun <V> get(key: Key<V>): V? = data.get(key)
 
     override fun compareTo(other: Interval): Int = lines.first - other.lines.first
+
+    fun getCellRange(editor: Editor): TextRange {
+      val document = editor.document
+      val startOffset = document.getLineNumber(document.getLineStartOffset(lines.first))
+      val endOffset = document.getLineNumber(document.getLineEndOffset(lines.last))
+      return TextRange(startOffset, endOffset)
+    }
+
+    fun getContentRange(editor: Editor): TextRange {
+      val contentLines = this.contentLines
+      val startOffset = editor.document.getLineStartOffset(contentLines.first)
+      val endOffset = editor.document.getLineEndOffset(contentLines.last)
+      return TextRange(startOffset, endOffset)
+    }
+
+    fun getContentText(editor: Editor): String {
+      val range = getContentRange(editor)
+      return editor.document.getText(range)
+    }
+
+    fun getCellText(editor: Editor): String {
+      val range = getCellRange(editor)
+      return editor.document.getText(range)
+    }
+
+
+    fun getTopMarker(document: Document): String? =
+      if (markers.hasTopLine) document.getLineText(lines.first) else null
+
+    fun getBottomMarker(document: Document): String? =
+      if (markers.hasBottomLine) document.getLineText(lines.last) else null
   }
 
   interface IntervalListener : EventListener {

notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookCellLinesContext.kt

@@ -1,100 +0,0 @@
-package com.intellij.notebooks.visualization
-
-import com.intellij.ide.IdeEventQueue
-import com.intellij.ide.impl.dataRules.GetDataRule
-import com.intellij.openapi.actionSystem.*
-import com.intellij.openapi.application.runReadAction
-import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.editor.ex.EditorGutterComponentEx
-import com.intellij.openapi.editor.impl.EditorComponentImpl
-import com.intellij.util.asSafely
-import com.intellij.util.containers.addIfNotNull
-import java.awt.Component
-import java.awt.event.MouseEvent
-import javax.swing.SwingUtilities
-
-/**
- * A list of editors and offsets inside them. Editors are ordered according to their conformity to the UI event and context,
- * and offsets represent the places in the editors closest to the place where the event happened.
- * The event and the context are extracted from the focused component, the mouse cursor position, and the caret position.
- */
-val EDITORS_WITH_OFFSETS_DATA_KEY: DataKey<List<Pair<Editor, Int>>> = DataKey.create("EDITORS_WITH_OFFSETS_DATA_KEY")
-
-private class NotebookCellLinesIntervalDataRule : GetDataRule {
-  override fun getData(dataProvider: DataProvider): NotebookCellLines.Interval? =
-    EDITORS_WITH_OFFSETS_DATA_KEY.getData(dataProvider)
-      ?.firstOrNull { (editor, _) ->
-        NotebookCellLinesProvider.get(editor.document) != null
-      }
-      ?.let { (editor, offset) ->
-        NotebookCellLines.get(editor).intervalsIterator(editor.document.getLineNumber(offset)).takeIf { it.hasNext() }?.next()
-      }
-}
-
-private class EditorsWithOffsetsDataRule : UiDataRule {
-  override fun uiDataSnapshot(sink: DataSink, snapshot: DataSnapshot) {
-    val contextComponent = snapshot[PlatformCoreDataKeys.CONTEXT_COMPONENT]
-    val editor = snapshot[PlatformDataKeys.EDITOR]
-    val result = runReadAction { getEditorsWithOffsets(editor, contextComponent) }
-    sink[EDITORS_WITH_OFFSETS_DATA_KEY] = result.takeIf(List<*>::isNotEmpty)
-  }
-}
-
-private fun getEditorsWithOffsets(editor: Editor?, contextComponent: Component?): List<Pair<Editor, Int>> {
-  // TODO Simplify. The code below is overcomplicated
-  val result = mutableListOf<Pair<Editor, Int>>()
-
-  // If the focused component is the editor, it's assumed that the current cell is the cell under the caret.
-  result.addIfNotNull(
-    contextComponent
-      ?.asSafely<EditorComponentImpl>()
-      ?.editor
-      ?.let { contextEditor ->
-        if (NotebookCellLinesProvider.get(contextEditor.document) != null) {
-          contextEditor to contextEditor.getOffsetFromCaretImpl()
-        } else null
-      })
-
-  // Otherwise, some component inside an editor can be focused. In that case it's assumed that the current cell is the cell closest
-  // to the focused component.
-  result.addIfNotNull(getOffsetInEditorByComponentHierarchy(contextComponent))
-
-  // When a user clicks on a gutter, it's the only focused component, and it's not connected to the editor. However, vertical offsets
-  // in the gutter can be juxtaposed to the editor.
-  if (contextComponent is EditorGutterComponentEx && editor != null) {
-    val mouseEvent = IdeEventQueue.getInstance().trueCurrentEvent as? MouseEvent
-    if (mouseEvent != null) {
-      val point = SwingUtilities.convertMouseEvent(mouseEvent.component, mouseEvent, contextComponent).point
-      result += editor to editor.logicalPositionToOffset(editor.xyToLogicalPosition(point))
-    }
-  }
-
-  // If the focused component is out of the notebook editor, there still can be other editors inside the required one.
-  // If some of such editors is treated as the current editor, the current cell is the cell closest to the current editor.
-  result.addIfNotNull(getOffsetInEditorByComponentHierarchy(editor?.contentComponent))
-
-  // When a user clicks on some toolbar on some menu component, it becomes the focused components. Usually, such components have an
-  // assigned editor. In that case it's assumed that the current cell is the cell under the caret.
-  if (editor != null && NotebookCellLinesProvider.get(editor.document) != null) {
-    result += editor to editor.getOffsetFromCaretImpl()
-  }
-
-  return result
-}
-
-private fun Editor.getOffsetFromCaretImpl(): Int =
-  caretModel.offset.coerceAtMost(document.textLength - 1).coerceAtLeast(0)
-
-private fun getOffsetInEditorByComponentHierarchy(component: Component?): Pair<Editor, Int>? =
-  generateSequence(component, Component::getParent)
-    .zipWithNext()
-    .mapNotNull { (child, parent) ->
-      if (parent is EditorComponentImpl) child to parent.editor
-      else null
-    }
-    .firstOrNull()
-    ?.let { (child, editor) ->
-      val point = SwingUtilities.convertPoint(child, 0, 0, editor.contentComponent)
-      editor to editor.logicalPositionToOffset(editor.xyToLogicalPosition(point))
-    }
-

notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookCellSelectionModelProvider.kt

@@ -3,8 +3,6 @@ package com.intellij.notebooks.visualization
 import com.intellij.lang.LanguageExtension
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.Key
-import com.intellij.util.concurrency.ThreadingAssertions
-import com.intellij.util.concurrency.annotations.RequiresEdt
 
 interface NotebookCellSelectionModelProvider {
   fun create(editor: Editor): NotebookCellSelectionModel
@@ -13,8 +11,7 @@ interface NotebookCellSelectionModelProvider {
 }
 
 val Editor.cellSelectionModel: NotebookCellSelectionModel?
-  @RequiresEdt get() {
-    ThreadingAssertions.assertEventDispatchThread()
+  get() {
     return key.get(this) ?: install(this)
   }
 

notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookIntervalPointer.kt

@@ -1,6 +1,5 @@
 package com.intellij.notebooks.visualization
 
-import com.intellij.openapi.actionSystem.DataKey
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
@@ -12,8 +11,6 @@ import com.intellij.util.concurrency.annotations.RequiresWriteLock
 import com.intellij.util.messages.Topic
 import java.util.*
 
-val NOTEBOOK_INTERVAL_POINTER_KEY = DataKey.create<NotebookIntervalPointer>("NOTEBOOK_INTERVAL_POINTER")
-
 /**
  * Pointer becomes invalid when code cell is removed.
  * It may become valid again when action is undone or redone.

notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContext.kt

@@ -0,0 +1,122 @@
+package com.intellij.notebooks.visualization.context
+
+import com.intellij.ide.IdeEventQueue
+import com.intellij.notebooks.visualization.NotebookCellLines
+import com.intellij.notebooks.visualization.NotebookCellLinesProvider
+import com.intellij.notebooks.visualization.cellSelectionModel
+import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.actionSystem.DataKey
+import com.intellij.openapi.actionSystem.PlatformCoreDataKeys
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.ex.EditorGutterComponentEx
+import com.intellij.openapi.editor.impl.EditorComponentImpl
+import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.util.asSafely
+import com.intellij.util.containers.addIfNotNull
+import java.awt.Component
+import java.awt.event.MouseEvent
+import javax.swing.SwingUtilities
+
+object NotebookDataContext {
+  val NOTEBOOK_CELL_LINES_INTERVAL = DataKey.create<NotebookCellLines.Interval>("NOTEBOOK_CELL_LINES_INTERVAL")
+
+  val DataContext.notebookEditor: EditorImpl?
+    get() {
+      val component = getData(PlatformCoreDataKeys.CONTEXT_COMPONENT) ?: return null
+      val editor = getData(PlatformCoreDataKeys.EDITOR)
+      val noteEditor = NotebookDataContextUtils.getCurrentEditor(editor, component) ?: return null
+      if (NotebookDataContextUtils.hasFocusedSearchReplaceComponent(noteEditor, component))
+        return null
+      return noteEditor
+    }
+
+  val DataContext.selectedCellIntervals: List<NotebookCellLines.Interval>?
+    get() {
+      val jupyterEditor = notebookEditor ?: return null
+      val selectionModel = jupyterEditor.cellSelectionModel ?: return null
+      return selectionModel.selectedCells
+    }
+
+
+  val DataContext.selectedCellInterval: NotebookCellLines.Interval?
+    get() {
+      val editor = notebookEditor ?: return null
+      val selectionModel = editor.cellSelectionModel ?: return null
+      return selectionModel.primarySelectedCell
+    }
+
+  val DataContext.hoveredInterval: NotebookCellLines.Interval?
+    get() {
+      val cached = getData(NOTEBOOK_CELL_LINES_INTERVAL)
+      if (cached != null)
+        return cached
+
+      val component = getData(PlatformCoreDataKeys.CONTEXT_COMPONENT)
+      val editor = notebookEditor ?: return null
+      val hoveredLine = NotebookDataContextUtils.getHoveredLine(editor, component) ?: return null
+      return NotebookCellLines.get(editor).getCell(hoveredLine)
+    }
+
+  val DataContext.hoveredOrSelectedInterval: NotebookCellLines.Interval?
+    get() = hoveredInterval ?: selectedCellInterval
+
+
+  private fun getEditorsWithOffsets(editor: Editor?, contextComponent: Component?): List<Pair<Editor, Int>> {
+    // TODO Simplify. The code below is overcomplicated
+    val result = mutableListOf<Pair<Editor, Int>>()
+
+    // If the focused component is the editor, it's assumed that the current cell is the cell under the caret.
+    result.addIfNotNull(
+      contextComponent
+        ?.asSafely<EditorComponentImpl>()
+        ?.editor
+        ?.let { contextEditor ->
+          if (NotebookCellLinesProvider.get(contextEditor.document) != null) {
+            contextEditor to contextEditor.getOffsetFromCaretImpl()
+          }
+          else null
+        })
+
+    // Otherwise, some component inside an editor can be focused. In that case it's assumed that the current cell is the cell closest
+    // to the focused component.
+    result.addIfNotNull(getOffsetInEditorByComponentHierarchy(contextComponent))
+
+    // When a user clicks on a gutter, it's the only focused component, and it's not connected to the editor. However, vertical offsets
+    // in the gutter can be juxtaposed to the editor.
+    if (contextComponent is EditorGutterComponentEx && editor != null) {
+      val mouseEvent = IdeEventQueue.getInstance().trueCurrentEvent as? MouseEvent
+      if (mouseEvent != null) {
+        val point = SwingUtilities.convertMouseEvent(mouseEvent.component, mouseEvent, contextComponent).point
+        result += editor to editor.logicalPositionToOffset(editor.xyToLogicalPosition(point))
+      }
+    }
+
+    // If the focused component is out of the notebook editor, there still can be other editors inside the required one.
+    // If some of such editors is treated as the current editor, the current cell is the cell closest to the current editor.
+    result.addIfNotNull(getOffsetInEditorByComponentHierarchy(editor?.contentComponent))
+
+    // When a user clicks on some toolbar on some menu component, it becomes the focused components. Usually, such components have an
+    // assigned editor. In that case it's assumed that the current cell is the cell under the caret.
+    if (editor != null && NotebookCellLinesProvider.get(editor.document) != null) {
+      result += editor to editor.getOffsetFromCaretImpl()
+    }
+
+    return result
+  }
+
+  private fun Editor.getOffsetFromCaretImpl(): Int =
+    caretModel.offset.coerceAtMost(document.textLength - 1).coerceAtLeast(0)
+
+  private fun getOffsetInEditorByComponentHierarchy(component: Component?): Pair<Editor, Int>? =
+    generateSequence(component, Component::getParent)
+      .zipWithNext()
+      .mapNotNull { (child, parent) ->
+        if (parent is EditorComponentImpl) child to parent.editor
+        else null
+      }
+      .firstOrNull()
+      ?.let { (child, editor) ->
+        val point = SwingUtilities.convertPoint(child, 0, 0, editor.contentComponent)
+        editor to editor.logicalPositionToOffset(editor.xyToLogicalPosition(point))
+      }
+}

notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContextUtils.kt

@@ -0,0 +1,85 @@
+package com.intellij.notebooks.visualization.context
+
+import com.intellij.find.SearchReplaceComponent
+import com.intellij.ide.IdeEventQueue
+import com.intellij.notebooks.visualization.NotebookCellLines
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.ex.EditorGutterComponentEx
+import com.intellij.openapi.editor.impl.EditorComponentImpl
+import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.util.asSafely
+import java.awt.Component
+import java.awt.event.MouseEvent
+import javax.swing.SwingUtilities
+
+object NotebookDataContextUtils {
+  fun getHoveredLine(editor: Editor?, contextComponent: Component?): Int? {
+    val noteEditor = editor as? EditorImpl ?: return null
+
+    // When a user clicks on a gutter, it's the only focused component, and it's not connected to the editor. However, vertical offsets
+    // in the gutter can be juxtaposed to the editor.
+    if (contextComponent is EditorGutterComponentEx) {
+      val mouseEditorOffset = getMouseEditorOffset(noteEditor, contextComponent)
+      if (mouseEditorOffset != null)
+        return mouseEditorOffset
+    }
+
+
+
+    val point = editor.contentComponent.mousePosition ?: return null
+    val logicalPosition = editor.xyToLogicalPosition(point)
+    return logicalPosition.line
+  }
+
+  fun getCurrentEditor(editor: Editor?, contextComponent: Component?): EditorImpl? {
+    val cachedEditor = editor?.takeIf { NotebookCellLines.hasSupport(it) }
+    if (cachedEditor != null) {
+      return cachedEditor as? EditorImpl
+    }
+
+    val componentSequence = generateSequence(contextComponent) {
+      it.parent
+    }
+
+    val noteEditor = componentSequence.firstNotNullOfOrNull { component ->
+      val editorComponentImpl = component as? EditorComponentImpl ?: return@firstNotNullOfOrNull null
+      val noteEditor = editorComponentImpl.editor
+      if (NotebookCellLines.hasSupport(noteEditor))
+        return@firstNotNullOfOrNull noteEditor
+
+      null
+    }
+    return noteEditor
+  }
+
+  fun getCaretLine(editor: EditorImpl): Int {
+    return editor.caretModel.logicalPosition.line
+  }
+
+  fun getRespectiveLineNumberInEditor(editor: Editor, component: Component): Int? {
+    val point = SwingUtilities.convertPoint(component, 0, component.height, editor.contentComponent)
+    val documentLineCount = editor.document.lineCount
+
+    if (point.y < 0)
+      return null
+
+    var prospectiveLineNumber = editor.xyToLogicalPosition(point).line
+    if (prospectiveLineNumber >= documentLineCount) {
+      prospectiveLineNumber = documentLineCount - 1
+    }
+    return prospectiveLineNumber
+  }
+
+
+  private fun getMouseEditorOffset(editor: Editor, contextComponent: Component?): Int? {
+    val mouseEvent = IdeEventQueue.getInstance().trueCurrentEvent as? MouseEvent ?: return null
+    val point = SwingUtilities.convertMouseEvent(mouseEvent.component, mouseEvent, contextComponent).point
+    return editor.logicalPositionToOffset(editor.xyToLogicalPosition(point))
+  }
+
+  fun hasFocusedSearchReplaceComponent(editor: Editor, contextComponent: Component): Boolean {
+    val searchReplaceComponent = editor.headerComponent.asSafely<SearchReplaceComponent>() ?: return false
+    return contextComponent === searchReplaceComponent.searchTextComponent ||
+           contextComponent === searchReplaceComponent.replaceTextComponent
+  }
+}

notebooks/visualization/src/com/intellij/notebooks/visualization/outputs/impl/NotebookOutputCollapseAllInSelectedCellsAction.kt

@@ -1,12 +1,19 @@
 package com.intellij.notebooks.visualization.outputs.impl
 
-import com.intellij.openapi.actionSystem.*
+import com.intellij.notebooks.visualization.NotebookCellInlayManager
+import com.intellij.notebooks.visualization.NotebookCellLines
+import com.intellij.notebooks.visualization.cellSelectionModel
+import com.intellij.notebooks.visualization.context.NotebookDataContext.notebookEditor
+import com.intellij.notebooks.visualization.context.NotebookDataContext.selectedCellIntervals
+import com.intellij.notebooks.visualization.ui.EditorCellOutput
+import com.intellij.notebooks.visualization.ui.NOTEBOOK_CELL_OUTPUT_DATA_KEY
+import com.intellij.openapi.actionSystem.ActionUpdateThread
+import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.actionSystem.PlatformDataKeys
+import com.intellij.openapi.actionSystem.ToggleAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.project.DumbAware
-import com.intellij.notebooks.visualization.*
-import com.intellij.notebooks.visualization.ui.EditorCellOutput
-import com.intellij.notebooks.visualization.ui.NOTEBOOK_CELL_OUTPUT_DATA_KEY
 
 internal class NotebookOutputCollapseAllAction private constructor() : ToggleAction(), DumbAware {
   override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.BGT
@@ -39,7 +46,7 @@ internal class NotebookOutputCollapseAllInSelectedCellsAction private constructo
     super.update(e)
     val editor = e.notebookEditor
     e.presentation.isEnabled = editor != null
-    e.presentation.isVisible = editor?.cellSelectionModel?.let { it.selectedCells.size > 1 } ?: false
+    e.presentation.isVisible = editor?.cellSelectionModel?.let { it.selectedCells.size > 1 } == true
   }
 
   override fun isSelected(e: AnActionEvent): Boolean =
@@ -87,7 +94,7 @@ internal class NotebookOutputCollapseSingleInCellAction private constructor() :
   }
 
   override fun isSelected(e: AnActionEvent): Boolean =
-    getExpectedComponent(e)?.collapsed ?: false
+    getExpectedComponent(e)?.collapsed == true
 
   override fun setSelected(e: AnActionEvent, state: Boolean) {
     getExpectedComponent(e)?.collapsed = state
@@ -98,9 +105,10 @@ internal class NotebookOutputCollapseSingleInCellAction private constructor() :
 }
 
 private fun getCollapsingComponents(e: AnActionEvent): List<EditorCellOutput> {
-  return e.dataContext.getData(NOTEBOOK_CELL_LINES_INTERVAL_DATA_KEY)?.let { interval ->
-    e.getData(PlatformDataKeys.EDITOR)?.let { getCollapsingComponents(it, interval) }
-  } ?: emptyList()
+  val selectedCellIntervals = e.dataContext.selectedCellIntervals ?: return emptyList()
+  val notebookEditor = e.dataContext.notebookEditor ?: return emptyList()
+
+  return selectedCellIntervals.flatMap { getCollapsingComponents(notebookEditor, it) }
 }
 
 private fun getCollapsingComponents(editor: Editor, interval: NotebookCellLines.Interval): List<EditorCellOutput> =

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

@@ -1,6 +1,13 @@
 package com.intellij.notebooks.visualization.ui
 
 import com.intellij.ide.DataManager
+import com.intellij.notebooks.ui.visualization.NotebookCodeCellBackgroundLineMarkerRenderer
+import com.intellij.notebooks.ui.visualization.NotebookLineMarkerRenderer
+import com.intellij.notebooks.ui.visualization.NotebookTextCellBackgroundLineMarkerRenderer
+import com.intellij.notebooks.ui.visualization.notebookAppearance
+import com.intellij.notebooks.visualization.*
+import com.intellij.notebooks.visualization.NotebookCellInlayController.InputFactory
+import com.intellij.notebooks.visualization.context.NotebookDataContext
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.DataProvider
@@ -17,12 +24,6 @@ import com.intellij.openapi.editor.markup.*
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.registry.Registry
 import com.intellij.util.asSafely
-import com.intellij.notebooks.ui.visualization.NotebookCodeCellBackgroundLineMarkerRenderer
-import com.intellij.notebooks.ui.visualization.NotebookLineMarkerRenderer
-import com.intellij.notebooks.ui.visualization.NotebookTextCellBackgroundLineMarkerRenderer
-import com.intellij.notebooks.ui.visualization.notebookAppearance
-import com.intellij.notebooks.visualization.*
-import com.intellij.notebooks.visualization.NotebookCellInlayController.InputFactory
 import java.awt.Graphics
 import java.awt.Graphics2D
 import java.awt.Point
@@ -399,7 +400,7 @@ class EditorCellView(
   ) : DataProvider {
     override fun getData(key: String): Any? =
       when (key) {
-        NOTEBOOK_CELL_LINES_INTERVAL_DATA_KEY.name -> intervalProvider()
+        NotebookDataContext.NOTEBOOK_CELL_LINES_INTERVAL.name -> intervalProvider()
         PlatformCoreDataKeys.CONTEXT_COMPONENT.name -> component
         PlatformDataKeys.EDITOR.name -> editor
         else -> null