diff --git a/notebooks/visualization/resources/intellij.notebooks.visualization.xml b/notebooks/visualization/resources/intellij.notebooks.visualization.xml
index f68088840da2..80122b26a081 100644
--- a/notebooks/visualization/resources/intellij.notebooks.visualization.xml
+++ b/notebooks/visualization/resources/intellij.notebooks.visualization.xml
@@ -102,6 +102,8 @@
     <notebookCellSelfManagedController
       implementation="com.intellij.notebooks.visualization.ui.cell.runCell.EditorCellRunGutterControllerFactory"/>
 
+    <notebookCellSelectionModelProvider language="Python"
+                                        implementationClass="com.intellij.notebooks.visualization.PythonCellSelectionModelProvider"/>
   </extensions>
 
   <extensions defaultExtensionNs="org.intellij.markdown">
diff --git a/notebooks/visualization/src/com/intellij/notebooks/visualization/CaretBasedCellSelectionModel.kt b/notebooks/visualization/src/com/intellij/notebooks/visualization/CaretBasedCellSelectionModel.kt
index 1ee59161d84c..aebf87035a61 100644
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/CaretBasedCellSelectionModel.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/CaretBasedCellSelectionModel.kt
@@ -1,3 +1,4 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.notebooks.visualization
 
 import com.intellij.openapi.application.ReadAction
@@ -7,6 +8,12 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.LogicalPosition
 import kotlin.math.min
 
+
+class PythonCellSelectionModelProvider : NotebookCellSelectionModelProvider {
+  override fun create(editor: Editor): NotebookCellSelectionModel =
+    CaretBasedCellSelectionModel(editor)
+}
+
 class CaretBasedCellSelectionModel(private val editor: Editor) : NotebookCellSelectionModel {
   override val primarySelectedCell: NotebookCellLines.Interval
     get() = editor.getCell(editor.caretModel.primaryCaret.logicalPosition.line)