PY-73153 [Jupyter] Add a number of collapsed cells after collapsed markdown heading

GitOrigin-RevId: 5c4f738f63c43f0d7ebc10ef5004c74e847dd67b

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/CustomFoldingEditorCellViewComponent.kt

@@ -1,18 +1,26 @@
 package org.jetbrains.plugins.notebooks.visualization.ui
 
+import com.intellij.codeInsight.hints.presentation.InlayPresentation
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.editor.CustomFoldRegion
 import com.intellij.openapi.editor.CustomFoldRegionRenderer
 import com.intellij.openapi.editor.ex.EditorEx
+import com.intellij.openapi.editor.impl.EditorGutterColor
+import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.editor.markup.GutterIconRenderer
 import com.intellij.openapi.editor.markup.TextAttributes
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.plugins.notebooks.visualization.UpdateContext
 import org.jetbrains.plugins.notebooks.visualization.ui.EditorEmbeddedComponentLayoutManager.CustomFoldingConstraint
+import java.awt.BorderLayout
+import java.awt.Dimension
+import java.awt.Graphics
 import java.awt.Graphics2D
 import java.awt.Rectangle
 import java.awt.geom.Rectangle2D
+import javax.swing.BoxLayout
 import javax.swing.JComponent
+import javax.swing.JPanel
 
 class CustomFoldingEditorCellViewComponent(
   internal val component: JComponent,
@@ -24,6 +32,15 @@ class CustomFoldingEditorCellViewComponent(
 
   private var gutterActionRenderer: ActionToGutterRendererAdapter? = null
 
+  private val bottomContainer = JPanel().apply {
+    layout = BoxLayout(this, BoxLayout.Y_AXIS)
+  }
+  private val mainComponent = JPanel().also {
+    it.layout = BorderLayout()
+    it.add(component, BorderLayout.CENTER)
+    it.add(bottomContainer, BorderLayout.SOUTH)
+  }
+
   @TestOnly
   fun getComponentForTest(): JComponent {
     return component
@@ -40,7 +57,7 @@ class CustomFoldingEditorCellViewComponent(
 
   private fun disposeFolding() {
     if (editor.isDisposed || foldingRegion?.isValid != true) return
-    editor.componentContainer.remove(component)
+    editor.componentContainer.remove(mainComponent)
     foldingRegion?.let {
       editor.foldingModel.runBatchFoldingOperation(
         {
@@ -74,7 +91,32 @@ class CustomFoldingEditorCellViewComponent(
         }
       }) ?: error("Failed to create folding region ${cell.interval.lines}")
       foldingRegion = fr
-      editor.componentContainer.add(component, CustomFoldingConstraint(fr, true))
+      editor.componentContainer.add(mainComponent, CustomFoldingConstraint(fr, true))
     }
   }
+
+  private val presentationToComponent = mutableMapOf<InlayPresentation, JComponent>()
+
+  override fun addInlayBelow(presentation: InlayPresentation) {
+    val inlayComponent = object : JComponent() {
+      override fun getPreferredSize(): Dimension? {
+        return Dimension(presentation.width, presentation.height)
+      }
+
+      override fun paint(g: Graphics) {
+        g as Graphics2D
+        val attributes = TextAttributes().apply {
+          backgroundColor = EditorGutterColor.getEditorGutterBackgroundColor(editor as EditorImpl, false)
+        }
+        presentation.paint(g, attributes)
+      }
+    }
+    inlayComponent.background = EditorGutterColor.getEditorGutterBackgroundColor(editor as EditorImpl, false)
+    presentationToComponent[presentation] = inlayComponent
+    bottomContainer.add(inlayComponent)
+  }
+
+  override fun removeInlayBelow(presentation: InlayPresentation) {
+    presentationToComponent.remove(presentation)?.let { bottomContainer.remove(it) }
+  }
 }

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellViewComponent.kt

@@ -1,5 +1,6 @@
 package org.jetbrains.plugins.notebooks.visualization.ui
 
+import com.intellij.codeInsight.hints.presentation.InlayPresentation
 import com.intellij.openapi.editor.Inlay
 import org.jetbrains.plugins.notebooks.visualization.UpdateContext
 import java.awt.Rectangle
@@ -48,4 +49,12 @@ abstract class EditorCellViewComponent {
   open fun doGetInlays(): Sequence<Inlay<*>> {
     return emptySequence()
   }
+
+  open fun addInlayBelow(presentation: InlayPresentation) {
+    throw UnsupportedOperationException("Operation is not supported")
+  }
+
+  open fun removeInlayBelow(presentation: InlayPresentation) {
+    throw UnsupportedOperationException("Operation is not supported")
+  }
 }

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/TextEditorCellViewComponent.kt

@@ -1,12 +1,17 @@
 package org.jetbrains.plugins.notebooks.visualization.ui
 
+import com.intellij.codeInsight.hints.presentation.InlayPresentation
+import com.intellij.codeInsight.hints.presentation.PresentationRenderer
 import com.intellij.openapi.actionSystem.AnAction
+import com.intellij.openapi.editor.Inlay
+import com.intellij.openapi.editor.InlayProperties
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.editor.impl.FoldingModelImpl
 import com.intellij.openapi.editor.markup.HighlighterLayer
 import com.intellij.openapi.editor.markup.HighlighterTargetArea
 import com.intellij.openapi.editor.markup.RangeHighlighter
 import com.intellij.openapi.editor.markup.TextAttributes
+import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.plugins.notebooks.visualization.NotebookCellLines
@@ -107,4 +112,22 @@ class TextEditorCellViewComponent(
     val offset = editor.document.getLineStartOffset(lines.first + 1)
     editor.caretModel.moveToOffset(offset)
   }
+
+  private val presentationToInlay = mutableMapOf<InlayPresentation, Inlay<*>>()
+
+  override fun addInlayBelow(presentation: InlayPresentation) {
+    editor.inlayModel.addBlockElement(
+      editor.document.getLineEndOffset(cell.interval.lines.last),
+      InlayProperties()
+        .showAbove(false)
+        .showWhenFolded(true),
+      PresentationRenderer(presentation)
+    )?.also { inlay ->
+      presentationToInlay[presentation] = inlay
+    }
+  }
+
+  override fun removeInlayBelow(presentation: InlayPresentation) {
+    presentationToInlay.remove(presentation)?.let { inlay -> Disposer.dispose(inlay) }
+  }
 }