PY-64378 Move code validation from full line to platform

GitOrigin-RevId: 3e31ea52588a3ee2dde035a149f2e4c45c368b8e
2026-03-22 15:19:59 +07:00 · 2023-11-28 13:32:55 +02:00
parent a75a94c36f
commit ebd8ebab42
25 changed files with 1135 additions and 3 deletions
--- a/platform/ml-impl/resources/META-INF/ml.xml
+++ b/platform/ml-impl/resources/META-INF/ml.xml
@@ -16,7 +16,11 @@
  </projectListeners>

  <extensionPoints>
-  <extensionPoint qualifiedName="com.intellij.platform.ml.impl.turboComplete.smartPipelineRunner"
-                  interface="com.intellij.platform.ml.impl.turboComplete.SmartPipelineRunner" dynamic="true"/>
+    <extensionPoint qualifiedName="com.intellij.platform.ml.impl.turboComplete.smartPipelineRunner"
+                    interface="com.intellij.platform.ml.impl.turboComplete.SmartPipelineRunner" dynamic="true"/>
+
+    <extensionPoint name="mlCompletionCorrectnessSupporter" beanClass="com.intellij.platform.ml.impl.correctness.MLCompletionCorrectnessSupporterEP" dynamic="true">
+      <with attribute="implementationClass" implements="com.intellij.platform.ml.impl.correctness.MLCompletionCorrectnessSupporter"/>
+    </extensionPoint>
  </extensionPoints>
 </idea-plugin>
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/MLCompletionCorrectnessSupporter.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/MLCompletionCorrectnessSupporter.kt
@@ -0,0 +1,19 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness
+
+import com.intellij.lang.Language
+import com.intellij.platform.ml.impl.correctness.autoimport.ImportFixer
+import com.intellij.platform.ml.impl.correctness.checker.CorrectnessChecker
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+interface MLCompletionCorrectnessSupporter {
+  val correctnessChecker: CorrectnessChecker
+  val importFixer: ImportFixer
+
+  companion object {
+    fun getInstance(language: Language): MLCompletionCorrectnessSupporter? {
+      return MLCompletionCorrectnessSupporterEP.EP_NAME.lazySequence().firstOrNull { it.language == language.id }?.instance
+    }
+  }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/MLCompletionCorrectnessSupporterBase.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/MLCompletionCorrectnessSupporterBase.kt
@@ -0,0 +1,12 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness
+
+import com.intellij.platform.ml.impl.correctness.autoimport.ImportFixer
+import com.intellij.platform.ml.impl.correctness.checker.CorrectnessCheckerBase
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+abstract class MLCompletionCorrectnessSupporterBase : MLCompletionCorrectnessSupporter {
+  override val correctnessChecker = CorrectnessCheckerBase()
+  override val importFixer: ImportFixer = ImportFixer.EMPTY
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/MLCompletionCorrectnessSupporterEP.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/MLCompletionCorrectnessSupporterEP.kt
@@ -0,0 +1,44 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness
+
+import com.intellij.openapi.extensions.CustomLoadingExtensionPointBean
+import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.extensions.RequiredElement
+import com.intellij.util.KeyedLazyInstance
+import com.intellij.util.xmlb.annotations.Attribute
+import org.jetbrains.annotations.ApiStatus
+import org.jetbrains.annotations.TestOnly
+
+@ApiStatus.Internal
+class MLCompletionCorrectnessSupporterEP : CustomLoadingExtensionPointBean<MLCompletionCorrectnessSupporter>,
+                                           KeyedLazyInstance<MLCompletionCorrectnessSupporter> {
+  @RequiredElement
+  @Attribute("language")
+  var language: String? = null
+
+  @RequiredElement
+  @Attribute("implementationClass")
+  var implementationClass: String? = null
+
+  @Suppress("unused")
+  constructor() : super()
+
+  @Suppress("unused")
+  @TestOnly
+  constructor(
+    language: String?,
+    implementationClass: String?,
+    instance: MLCompletionCorrectnessSupporter,
+  ) : super(instance) {
+    this.language = language
+    this.implementationClass = implementationClass
+  }
+
+  override fun getImplementationClassName(): String? = this.implementationClass
+  override fun getKey(): String = language!!
+
+  companion object {
+    val EP_NAME: ExtensionPointName<MLCompletionCorrectnessSupporterEP> = ExtensionPointName.create(
+      "com.intellij.mlCompletionCorrectnessSupporter")
+  }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/autoimport/ImportFixer.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/autoimport/ImportFixer.kt
@@ -0,0 +1,48 @@
+package com.intellij.platform.ml.impl.correctness.autoimport
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.readAction
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.progress.blockingContextToIndicator
+import com.intellij.openapi.project.DumbService
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiFile
+import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
+import com.intellij.util.concurrency.annotations.RequiresBlockingContext
+import com.intellij.util.concurrency.annotations.RequiresReadLock
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+interface ImportFixer {
+  @RequiresReadLock
+  @RequiresBackgroundThread
+  @RequiresBlockingContext
+  fun runAutoImport(file: PsiFile, editor: Editor, suggestionRange: TextRange)
+
+  object EMPTY : ImportFixer {
+    override fun runAutoImport(file: PsiFile, editor: Editor, suggestionRange: TextRange) {}
+  }
+}
+
+@RequiresBlockingContext
+@ApiStatus.Internal
+fun ImportFixer.runAutoImportAsync(scope: CoroutineScope, file: PsiFile, editor: Editor, suggestionRange: TextRange) {
+  val autoImportAction = {
+    if (!DumbService.getInstance(file.project).isDumb) {
+      runAutoImport(file, editor, suggestionRange)
+    }
+  }
+
+  if (ApplicationManager.getApplication().isUnitTestMode) {
+    autoImportAction()
+  }
+  else {
+    scope.launch {
+      readAction {
+        blockingContextToIndicator(autoImportAction)
+      }
+    }
+  }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/autoimport/InspectionBasedImportFixer.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/autoimport/InspectionBasedImportFixer.kt
@@ -0,0 +1,67 @@
+package com.intellij.platform.ml.impl.correctness.autoimport
+
+import com.intellij.codeInsight.daemon.impl.DaemonProgressIndicator
+import com.intellij.codeInspection.InspectionEngine
+import com.intellij.codeInspection.LocalInspectionTool
+import com.intellij.codeInspection.LocalQuickFixOnPsiElement
+import com.intellij.codeInspection.ProblemDescriptor
+import com.intellij.codeInspection.ex.LocalInspectionToolWrapper
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.ModalityState
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.SyntaxTraverser
+import com.intellij.util.PairProcessor
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+abstract class InspectionBasedImportFixer : ImportFixer {
+
+  protected abstract fun getAutoImportInspections(element: PsiElement?): List<LocalInspectionTool>
+  protected abstract fun filterApplicableFixes(fixes: List<LocalQuickFixOnPsiElement>): List<LocalQuickFixOnPsiElement>
+  override fun runAutoImport(file: PsiFile, editor: Editor, suggestionRange: TextRange) {
+    val elements = SyntaxTraverser.psiTraverser(file)
+      .onRange(suggestionRange)
+      .toList()
+    val indicator = when (ApplicationManager.getApplication().isUnitTestMode) {
+      false -> ProgressManager.getInstance().progressIndicator
+      true -> DaemonProgressIndicator()
+    }
+
+    val fixes = InspectionEngine.inspectElements(
+      getAutoImportInspections(file).map { LocalInspectionToolWrapper(it) },
+      file,
+      file.textRange,
+      true,
+      true,
+      indicator,
+      elements,
+      PairProcessor.alwaysTrue()
+    ).values.flatMap { problemDescriptors ->
+      problemDescriptors.flatMap { it.fixes.orEmpty().toList() }
+    }.filterIsInstance<LocalQuickFixOnPsiElement>()
+
+    applyFixes(editor, filterApplicableFixes(fixes))
+  }
+
+  fun areFixableByAutoImport(problems: List<ProblemDescriptor>): Boolean {
+    return problems.all {
+      val fixes = it.fixes.orEmpty().filterIsInstance<LocalQuickFixOnPsiElement>()
+      filterApplicableFixes(fixes).isNotEmpty()
+    }
+  }
+
+  open fun applyFixes(editor: Editor, fixes: List<LocalQuickFixOnPsiElement>) {
+    val fixToApply = fixes.firstOrNull() ?: return // To avoid layering of some import popups on others
+    val lastModified = editor.document.modificationStamp
+    fun action() = ApplicationManager.getApplication().runWriteAction {
+      fixToApply.applyFix()
+    }
+    ApplicationManager.getApplication().invokeLater(::action, ModalityState.defaultModalityState()) {
+      editor.document.modificationStamp != lastModified
+    }
+  }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CodeCorrectness.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CodeCorrectness.kt
@@ -0,0 +1,14 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness.checker
+
+import org.jetbrains.annotations.ApiStatus
+
+
+@ApiStatus.Internal
+data class CodeCorrectness(val syntaxState: ErrorsState, val semanticState: ErrorsState) {
+  fun allErrors(): List<CorrectnessError> = syntaxState.errors() + semanticState.errors()
+
+  companion object {
+    fun empty(): CodeCorrectness = CodeCorrectnessBuilder().build()
+  }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CodeCorrectnessBuilder.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CodeCorrectnessBuilder.kt
@@ -0,0 +1,40 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness.checker
+
+import com.intellij.platform.ml.impl.correctness.checker.ErrorsState.Unknown.UnknownReason
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+class CodeCorrectnessBuilder {
+  private var syntaxState: ErrorsState = ErrorsState.Unknown(UnknownReason.NOT_STARTED)
+  private var semanticState: ErrorsState = ErrorsState.Unknown(UnknownReason.NOT_STARTED)
+  fun syntaxCorrectness(block: () -> List<CorrectnessError>) {
+    syntaxState = ErrorsState.Unknown(UnknownReason.IN_PROGRESS)
+    val errors = block()
+    syntaxState = selectState(errors)
+  }
+
+  fun semanticCorrectness(block: () -> List<CorrectnessError>) {
+    semanticState = ErrorsState.Unknown(UnknownReason.IN_PROGRESS)
+    val errors = block()
+    semanticState = selectState(errors)
+  }
+
+  private fun selectState(errors: List<CorrectnessError>): ErrorsState = when {
+    errors.isEmpty() -> ErrorsState.Correct
+    else -> ErrorsState.Incorrect(errors)
+  }
+
+  fun timeLimitExceeded() {
+    if ((syntaxState as? ErrorsState.Unknown)?.reason == UnknownReason.IN_PROGRESS) {
+      syntaxState = ErrorsState.Unknown(UnknownReason.TIME_LIMIT_EXCEEDED)
+    }
+    if ((semanticState as? ErrorsState.Unknown)?.reason == UnknownReason.IN_PROGRESS) {
+      semanticState = ErrorsState.Unknown(UnknownReason.TIME_LIMIT_EXCEEDED)
+    }
+  }
+
+  fun build(): CodeCorrectness {
+    return CodeCorrectness(syntaxState, semanticState)
+  }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CorrectnessChecker.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CorrectnessChecker.kt
@@ -0,0 +1,23 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness.checker
+
+import com.intellij.psi.PsiFile
+import com.intellij.util.concurrency.annotations.RequiresBlockingContext
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+interface CorrectnessChecker {
+  @RequiresBlockingContext
+  @ApiStatus.ScheduledForRemoval
+  @Deprecated("Do not use it")
+  fun checkSyntax(file: PsiFile,
+                  suggestion: String,
+                  offset: Int,
+                  prefix: String, ignoreSyntaxErrorsBeforeSuggestionLen: Int): List<CorrectnessError>
+
+  @RequiresBlockingContext
+  fun checkSemantic(file: PsiFile,
+                    suggestion: String,
+                    offset: Int,
+                    prefix: String): List<CorrectnessError>
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CorrectnessCheckerBase.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CorrectnessCheckerBase.kt
@@ -0,0 +1,78 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness.checker
+
+import com.intellij.codeInspection.InspectionEngine
+import com.intellij.openapi.diagnostic.thisLogger
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiFile
+import com.intellij.psi.SyntaxTraverser
+import com.intellij.util.PairProcessor
+import com.intellij.platform.ml.impl.correctness.finalizer.SuggestionFinalizer
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+open class CorrectnessCheckerBase(open val semanticCheckers: List<SemanticChecker> = emptyList()) : CorrectnessChecker {
+  @Suppress("PropertyName")
+  protected val LOG = thisLogger()
+
+  protected open val suggestionFinalizer: SuggestionFinalizer? = null
+
+  final override fun checkSyntax(file: PsiFile,
+                                 suggestion: String,
+                                 offset: Int,
+                                 prefix: String,
+                                 ignoreSyntaxErrorsBeforeSuggestionLen: Int): List<CorrectnessError> {
+    // todo: consider using length in leaves instead of plain offset
+    val isSyntaxCorrect = suggestionFinalizer
+                            ?.getFinalization(file, suggestion, offset, prefix)
+                            ?.hasNoErrorsStartingFrom(offset - ignoreSyntaxErrorsBeforeSuggestionLen) ?: true
+    return if (isSyntaxCorrect) {
+      emptyList()
+    }
+    else {
+      listOf(CorrectnessError(TextRange.EMPTY_RANGE, Severity.CRITICAL)) // todo specify error location
+    }
+  }
+
+  protected open fun buildPsiForSemanticChecks(file: PsiFile, suggestion: String, offset: Int, prefix: String): PsiFile {
+    return file
+  }
+
+  private val toolWrappers = semanticCheckers.filterIsInstance<InspectionBasedSemanticChecker>()
+    .map { it.toolWrapper }
+
+  private val toolNameToSemanticChecker = semanticCheckers.filterIsInstance<InspectionBasedSemanticChecker>()
+    .associateBy { it.toolWrapper.id }
+
+  private val customSemanticCheckers = semanticCheckers.filterIsInstance<CustomSemanticChecker>()
+
+  final override fun checkSemantic(file: PsiFile, suggestion: String, offset: Int, prefix: String): List<CorrectnessError> {
+    if (semanticCheckers.isEmpty()) {
+      return emptyList()
+    }
+    val fullPsi = buildPsiForSemanticChecks(file, suggestion, offset, prefix)
+
+    val range = TextRange(offset - prefix.length, offset + suggestion.length - prefix.length)
+
+    val elements = SyntaxTraverser.psiTraverser(fullPsi)
+      .onRange(range)
+      .toList()
+
+    return InspectionEngine.inspectElements(
+      toolWrappers,
+      fullPsi,
+      fullPsi.textRange,
+      true,
+      true,
+      ProgressManager.getInstance().progressIndicator,
+      elements,
+      PairProcessor.alwaysTrue()
+    ).flatMap {
+      val semanticChecker = checkNotNull(toolNameToSemanticChecker[it.key.id])
+      semanticChecker.convertInspectionsResults(file, it.value, offset, prefix, suggestion)
+    } + customSemanticCheckers.flatMap { analyzer ->
+      elements.flatMap { element -> analyzer.findErrors(file, element, offset, prefix, suggestion) }
+    }
+  }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CorrectnessError.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/CorrectnessError.kt
@@ -0,0 +1,8 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness.checker
+
+import com.intellij.openapi.util.TextRange
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+data class CorrectnessError(val location: TextRange, val severity: Severity)
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/ErrorsState.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/ErrorsState.kt
@@ -0,0 +1,35 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness.checker
+
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+sealed class ErrorsState {
+  object Correct : ErrorsState() {
+    override fun toString(): String = "Correct"
+
+    override fun errors() = emptyList<CorrectnessError>()
+  }
+
+  data class Incorrect(val errors: List<CorrectnessError>) : ErrorsState() {
+    init {
+      require(errors.isNotEmpty())
+    }
+
+    override fun errors() = errors
+  }
+
+  data class Unknown(val reason: UnknownReason) : ErrorsState() {
+    enum class UnknownReason {
+      TIME_LIMIT_EXCEEDED,
+      NOT_STARTED,
+      IN_PROGRESS
+    }
+
+    override fun errors() = emptyList<CorrectnessError>()
+  }
+
+  abstract fun errors(): List<CorrectnessError>
+
+  fun List<CorrectnessError>.hasCriticalErrors(): Boolean = any { it.severity == Severity.CRITICAL }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/SemanticChecker.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/SemanticChecker.kt
@@ -0,0 +1,52 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness.checker
+
+import com.intellij.codeInspection.LocalInspectionTool
+import com.intellij.codeInspection.ProblemDescriptor
+import com.intellij.codeInspection.ex.LocalInspectionToolWrapper
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+sealed interface SemanticChecker {
+  fun getLocationInSuggestion(errorRangeInFile: TextRange, offset: Int, prefix: String, suggestion: String): TextRange? {
+    val shift = offset - prefix.length
+    val suggestionLocationInFile = TextRange(0, suggestion.length).shiftRight(shift)
+    if (!suggestionLocationInFile.intersects(errorRangeInFile)) {
+      return null
+    }
+    return suggestionLocationInFile.intersection(errorRangeInFile)!!.shiftLeft(shift)
+  }
+}
+
+@ApiStatus.Internal
+abstract class InspectionBasedSemanticChecker(localInspectionTool: LocalInspectionTool) : SemanticChecker {
+  abstract fun convertInspectionsResults(
+    originalPsi: PsiFile,
+    problemDescriptors: List<ProblemDescriptor>,
+    offset: Int,
+    prefix: String,
+    suggestion: String
+  ): List<CorrectnessError>
+
+  val toolWrapper: LocalInspectionToolWrapper = LocalInspectionToolWrapper(localInspectionTool)
+
+  protected fun getLocationInSuggestion(problemDescriptor: ProblemDescriptor, offset: Int, prefix: String, suggestion: String): TextRange? =
+    getLocationInSuggestion(getErrorRangeInFile(problemDescriptor), offset, prefix, suggestion)
+
+  protected fun getErrorRangeInFile(problemDescriptor: ProblemDescriptor): TextRange {
+    val rangeInElement = problemDescriptor.textRangeInElement ?: TextRange(0, problemDescriptor.psiElement.textLength)
+    return rangeInElement.shiftRight(problemDescriptor.psiElement.textRange.startOffset)
+  }
+}
+
+@ApiStatus.Internal
+abstract class CustomSemanticChecker : SemanticChecker {
+  abstract fun findErrors(originalPsi: PsiFile,
+                          element: PsiElement,
+                          offset: Int,
+                          prefix: String,
+                          suggestion: String): List<CorrectnessError>
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/Severity.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/checker/Severity.kt
@@ -0,0 +1,10 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.ml.impl.correctness.checker
+
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+enum class Severity {
+  CRITICAL,
+  ACCEPTABLE,
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/finalizer/FinalizedFile.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/finalizer/FinalizedFile.kt
@@ -0,0 +1,29 @@
+package com.intellij.platform.ml.impl.correctness.finalizer
+
+import com.intellij.psi.PsiErrorElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.SyntaxTraverser
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.ScheduledForRemoval
+@Deprecated("Do not use it")
+class FinalizedFile(finalizedPsi: PsiFile) {
+  private val errorElementsRanges = SyntaxTraverser.psiTraverser(finalizedPsi)
+    .filter(PsiErrorElement::class.java)
+    .map { it.textRange }
+    .toList()
+
+  /**
+   * This function checks that the finalized file does not contain errors starting with [offset] or later.
+   *
+   * Since when finalizing a file, we only add code after suggestion,
+   * if you run it with [offset] = 0
+   * it cannot return true if suggestion is syntactically incorrect.
+   * This is an important property of this function, it must be preserved.
+   *
+   * Of course, if you run this function with a large offset, it will always say that there are no errors.
+   * So you should run it with an offset less than the offset of the completion call.
+   * But how much less depends on how many errors you need to ignore before completion.
+   */
+  fun hasNoErrorsStartingFrom(offset: Int): Boolean = errorElementsRanges.none { it.startOffset >= offset }
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/finalizer/SuggestionFinalizer.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/finalizer/SuggestionFinalizer.kt
@@ -0,0 +1,33 @@
+package com.intellij.platform.ml.impl.correctness.finalizer
+
+import com.intellij.platform.ml.impl.correctness.finalizer.FinalizedFile
+import com.intellij.psi.PsiFile
+import org.jetbrains.annotations.ApiStatus.ScheduledForRemoval
+
+@ScheduledForRemoval
+@Deprecated("Do not use it")
+interface SuggestionFinalizer {
+  /**
+   * This function takes the text up to the caret, inserts the [suggestion] into it
+   * and tries to finalize the file so that the inserted [suggestion] does not create syntax errors
+   * (due to unfinished code in the [suggestion]).
+   *
+   * Example:
+   * If we have a context
+   * ```
+   * for i
+   * ```
+   * And suggestion
+   * ```
+   * in range(
+   * ```
+   * Then the [FinalizedFile] likely will be
+   * ```
+   * for i in range():
+   *  pass
+   * ```
+   *
+   * Please see [FinalizedFile] if you are going to change the semantics of this function.
+   */
+  fun getFinalization(originalPsi: PsiFile, suggestion: String, offset: Int, prefix: String): FinalizedFile
+}
--- a/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/finalizer/SuggestionFinalizerBase.kt
+++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/correctness/finalizer/SuggestionFinalizerBase.kt
@@ -0,0 +1,39 @@
+package com.intellij.platform.ml.impl.correctness.finalizer
+
+import com.intellij.lang.Language
+import com.intellij.openapi.application.runReadAction
+import com.intellij.psi.*
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.ScheduledForRemoval
+@Deprecated("Do not use it")
+abstract class SuggestionFinalizerBase(val language: Language) : SuggestionFinalizer {
+
+  /**
+   * Returns the order in which the elements will be traversed to finalize.
+   */
+  protected abstract fun getTraverseOrder(insertedElement: PsiElement): Sequence<PsiElement>
+
+  /**
+   * Returns the finalization for the element.
+   */
+  protected abstract fun getFinalizationCandidate(element: PsiElement): String?
+
+
+  /**
+   * This implementation traverse tree with the inserted [suggestion] according to [getTraverseOrder]
+   * and appends text from [getFinalizationCandidate].
+   */
+  final override fun getFinalization(originalPsi: PsiFile, suggestion: String, offset: Int, prefix: String): FinalizedFile = runReadAction {
+    require(suggestion.isNotBlank())
+    val psi = PsiFileFactory.getInstance(originalPsi.project)
+      .createFileFromText(language, originalPsi.text.take(offset - prefix.length) + suggestion)
+    val insertedElement = psi.findElementAt(psi.textLength - 1 - suggestion.takeLastWhile { it == ' ' }.length)!!
+    val finalization = getTraverseOrder(insertedElement).joinToString(separator = "") {
+      getFinalizationCandidate(it).orEmpty()
+    }
+    val finalizedText = psi.text + finalization
+    val finalizedPsi = PsiFileFactory.getInstance(originalPsi.project).createFileFromText(language, finalizedText)
+    FinalizedFile(finalizedPsi)
+  }
+}
--- a/python/intellij.python.community.impl.iml
+++ b/python/intellij.python.community.impl.iml
@@ -111,7 +111,6 @@
    <orderEntry type="library" name="kotlin-stdlib" level="project" />
    <orderEntry type="library" name="caffeine" level="project" />
    <orderEntry type="module" module-name="intellij.platform.util.jdom" />
-    <orderEntry type="module" module-name="intellij.platform.ml" />
    <orderEntry type="module" module-name="intellij.platform.extensions" />
    <orderEntry type="library" name="kotlinx-serialization-json" level="project" />
    <orderEntry type="library" name="kotlinx-serialization-core" level="project" />
@@ -143,5 +142,7 @@
    </orderEntry>
    <orderEntry type="module" module-name="intellij.platform.backend.observation" />
    <orderEntry type="module" module-name="intellij.terminal" />
+    <orderEntry type="module" module-name="intellij.platform.ml.impl" />
+    <orderEntry type="module" module-name="intellij.platform.ml" />
  </component>
 </module>
--- a/python/src/META-INF/python-core-common.xml
+++ b/python/src/META-INF/python-core-common.xml
@@ -264,6 +264,8 @@
                                   implementationClass="com.jetbrains.python.codeInsight.mlcompletion.PyContextFeatureProvider"/>
    <completion.ml.elementFeatures language="Python"
                                   implementationClass="com.jetbrains.python.codeInsight.mlcompletion.PyElementFeatureProvider"/>
+    <mlCompletionCorrectnessSupporter language="Python"
+                                      implementationClass="com.jetbrains.python.codeInsight.mlcompletion.correctness.PythonMLCompletionCorrectnessSupporter"/>
    <completion.confidence language="Python" implementationClass="com.jetbrains.python.codeInsight.completion.PyCompletionConfidence"/>
    <completion.ml.model implementation="com.jetbrains.python.codeInsight.mlcompletion.PythonMLRankingProvider"/>
    <typedHandler implementation="com.jetbrains.python.console.completion.PythonConsoleAutopopupBlockingHandler" id="pydevBlockAutoPopup"
--- a/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/PythonMLCompletionCorrectnessSupporter.kt
+++ b/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/PythonMLCompletionCorrectnessSupporter.kt
@@ -0,0 +1,12 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.codeInsight.mlcompletion.correctness
+
+import com.intellij.platform.ml.impl.correctness.MLCompletionCorrectnessSupporterBase
+import com.jetbrains.python.codeInsight.mlcompletion.correctness.autoimport.PythonImportFixer
+import com.jetbrains.python.codeInsight.mlcompletion.correctness.checker.PythonCorrectnessChecker
+
+class PythonMLCompletionCorrectnessSupporter : MLCompletionCorrectnessSupporterBase() {
+  override val correctnessChecker: PythonCorrectnessChecker = PythonCorrectnessChecker()
+
+  override val importFixer: PythonImportFixer = PythonImportFixer()
+}
--- a/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/autoimport/PythonImportFixer.kt
+++ b/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/autoimport/PythonImportFixer.kt
@@ -0,0 +1,29 @@
+package com.jetbrains.python.codeInsight.mlcompletion.correctness.autoimport
+
+import com.intellij.codeInspection.LocalQuickFixOnPsiElement
+import com.intellij.platform.ml.impl.correctness.autoimport.InspectionBasedImportFixer
+import com.intellij.psi.PsiDirectory
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.elementType
+import com.jetbrains.python.codeInsight.imports.AutoImportQuickFix
+import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferencesInspection
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyFileElementType
+
+class PythonImportFixer : InspectionBasedImportFixer() {
+  override fun getAutoImportInspections(element: PsiElement?) =
+    PyUnresolvedReferencesInspection.getInstance(element)?.let { listOf(it) } ?: listOf()
+
+  override fun filterApplicableFixes(fixes: List<LocalQuickFixOnPsiElement>): List<AutoImportQuickFix> {
+    val fixesToApply = fixes.filterIsInstance<AutoImportQuickFix>()
+      .filter { it.isAvailable }
+      .filter {
+        it.candidates.any {
+          val importable = it.importable
+          // Check that an importing element is a module or class
+          importable is PyClass || importable is PsiDirectory || importable.elementType is PyFileElementType
+        }
+      }
+    return fixesToApply
+  }
+}
--- a/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/checker/PythonCorrectnessChecker.kt
+++ b/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/checker/PythonCorrectnessChecker.kt
@@ -0,0 +1,26 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.codeInsight.mlcompletion.correctness.checker
+
+import com.intellij.platform.ml.impl.correctness.checker.CorrectnessCheckerBase
+import com.intellij.openapi.util.io.FileUtil
+import com.intellij.psi.PsiFile
+import com.jetbrains.python.codeInsight.mlcompletion.correctness.finalizer.PythonSuggestionFinalizer
+import com.jetbrains.python.psi.impl.PyExpressionCodeFragmentImpl
+
+class PythonCorrectnessChecker : CorrectnessCheckerBase(listOf(
+  PyUnresolvedReferencesSemanticChecker,
+  PyCallingNonCallableSemanticChecker,
+  PyArgumentListSemanticChecker,
+  PyRedeclarationSemanticChecker,
+  PyAssignmentToLibraryScopeSemanticChecker,
+)) {
+  override val suggestionFinalizer = PythonSuggestionFinalizer()
+  override fun buildPsiForSemanticChecks(file: PsiFile, suggestion: String, offset: Int, prefix: String): PsiFile {
+    return PyExpressionCodeFragmentImpl(
+      file.project,
+      FileUtil.getNameWithoutExtension(file.name) + ".py",
+      file.text.let { it.take(offset - prefix.length) + suggestion + " " + it.drop(offset) },
+      true
+    ).apply { context = file }
+  }
+}
--- a/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/checker/PythonCustomSemanticCheckers.kt
+++ b/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/checker/PythonCustomSemanticCheckers.kt
@@ -0,0 +1,36 @@
+package com.jetbrains.python.codeInsight.mlcompletion.correctness.checker
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.search.ProjectScope
+import com.jetbrains.python.psi.PyAssignmentStatement
+import com.jetbrains.python.psi.PyTargetExpression
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.TypeEvalContext
+import com.intellij.platform.ml.impl.correctness.checker.CorrectnessError
+import com.intellij.platform.ml.impl.correctness.checker.CustomSemanticChecker
+import com.intellij.platform.ml.impl.correctness.checker.Severity
+
+object PyAssignmentToLibraryScopeSemanticChecker : CustomSemanticChecker() {
+  override fun findErrors(originalPsi: PsiFile,
+                          element: PsiElement,
+                          offset: Int,
+                          prefix: String,
+                          suggestion: String): List<CorrectnessError> {
+    val assignment = element as? PyAssignmentStatement ?: return emptyList()
+    val typeEvalContext = TypeEvalContext.codeAnalysis(element.project, originalPsi)
+    val resolveContext = PyResolveContext.defaultContext(typeEvalContext)
+    val librariesScope = ProjectScope.getLibrariesScope(element.project)
+    return assignment.targets
+      .filterIsInstance<PyTargetExpression>()
+      .filter {
+        val declarationElement = it.getReference(resolveContext).resolve() ?: return@filter false
+        val virtualFile = declarationElement.containingFile?.virtualFile ?: return@filter false
+        librariesScope.contains(virtualFile)
+      }.mapNotNull {
+        val location = getLocationInSuggestion(it.textRange, offset, prefix, suggestion) ?: return@mapNotNull null
+        CorrectnessError(location, Severity.CRITICAL)
+      }
+  }
+
+}
--- a/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/checker/PythonInspectionBasedSemanticCheckers.kt
+++ b/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/checker/PythonInspectionBasedSemanticCheckers.kt
@@ -0,0 +1,117 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.codeInsight.mlcompletion.correctness.checker
+
+import com.intellij.codeInspection.ProblemDescriptor
+import com.intellij.codeInspection.ProblemHighlightType
+import com.intellij.platform.ml.impl.correctness.MLCompletionCorrectnessSupporter
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.PsiFile
+import com.intellij.psi.util.parentOfType
+import com.intellij.refactoring.suggested.startOffset
+import com.jetbrains.python.PyTokenTypes
+import com.jetbrains.python.PythonLanguage
+import com.jetbrains.python.codeInsight.mlcompletion.PyMlCompletionHelpers
+import com.jetbrains.python.inspections.PyArgumentListInspection
+import com.jetbrains.python.inspections.PyCallingNonCallableInspection
+import com.jetbrains.python.inspections.PyRedeclarationInspection
+import com.jetbrains.python.inspections.unresolvedReference.PyUnresolvedReferencesInspection
+import com.jetbrains.python.psi.PyFromImportStatement
+import com.jetbrains.python.psi.PyImportStatement
+import com.jetbrains.python.psi.PyImportStatementBase
+import com.intellij.platform.ml.impl.correctness.checker.CorrectnessError
+import com.intellij.platform.ml.impl.correctness.checker.InspectionBasedSemanticChecker
+import com.intellij.platform.ml.impl.correctness.checker.Severity
+import com.jetbrains.python.codeInsight.mlcompletion.correctness.PythonMLCompletionCorrectnessSupporter
+
+object PyUnresolvedReferencesSemanticChecker : InspectionBasedSemanticChecker(PyUnresolvedReferencesInspection()) {
+  override fun convertInspectionsResults(originalPsi: PsiFile,
+                                         problemDescriptors: List<ProblemDescriptor>,
+                                         offset: Int,
+                                         prefix: String,
+                                         suggestion: String): List<CorrectnessError> =
+    problemDescriptors.filterErrorsInsideUnresolvedWellKnownImports().mapNotNull { problemDescriptor ->
+      val severity = getErrorSeverity(problemDescriptor)
+      val location = getLocationInSuggestion(problemDescriptor, offset, prefix, suggestion) ?: return@mapNotNull null
+      CorrectnessError(location, severity)
+  }
+
+  private fun List<ProblemDescriptor>.filterErrorsInsideUnresolvedWellKnownImports(): List<ProblemDescriptor> {
+    val unresolvedWellKnownProblems = filter {
+      val fromStatement = it.psiElement.parentOfType<PyFromImportStatement>() ?: return@filter false
+      val importSource = fromStatement.importSource ?: return@filter false
+      getErrorRangeInFile(it) in importSource.textRange &&
+      it.psiElement.text in PyMlCompletionHelpers.importPopularity
+    } + filter {
+      val importStatement = it.psiElement.parentOfType<PyImportStatement>() ?: return@filter false
+      importStatement.importElements.all { element ->
+        val firstName = element.importedQName?.components?.firstOrNull() ?: return@all true
+        getErrorRangeInFile(it) in TextRange(0, firstName.length).shiftRight(element.startOffset) &&
+        firstName in PyMlCompletionHelpers.importPopularity
+      }
+    }
+    val ignoreStartOffsets = unresolvedWellKnownProblems.mapNotNull {
+      it.psiElement.parentOfType<PyImportStatementBase>()?.startOffset
+    }.toSet()
+    return filter {
+      val importStatement = it.psiElement.parentOfType<PyImportStatementBase>() ?: return@filter true
+      importStatement.startOffset !in ignoreStartOffsets
+    }
+  }
+
+  private fun getErrorSeverity(problemDescriptor: ProblemDescriptor): Severity {
+    problemDescriptor.psiElement.parentOfType<PyImportStatementBase>()?.let {
+      // errors inside import are critical
+      return Severity.CRITICAL
+    }
+
+    val importFixer = (MLCompletionCorrectnessSupporter.getInstance(PythonLanguage.INSTANCE) as PythonMLCompletionCorrectnessSupporter).importFixer
+    return if (importFixer.areFixableByAutoImport(listOf(problemDescriptor))) {
+      Severity.ACCEPTABLE
+    } else {
+      Severity.CRITICAL
+    }
+  }
+}
+
+object PyCallingNonCallableSemanticChecker : InspectionBasedSemanticChecker(PyCallingNonCallableInspection()) {
+  override fun convertInspectionsResults(originalPsi: PsiFile,
+                                         problemDescriptors: List<ProblemDescriptor>,
+                                         offset: Int,
+                                         prefix: String,
+                                         suggestion: String): List<CorrectnessError> =
+    problemDescriptors.mapNotNull { problemDescriptor ->
+      val location = getLocationInSuggestion(problemDescriptor, offset, prefix, suggestion) ?: return@mapNotNull null
+      CorrectnessError(location, Severity.CRITICAL)
+  }
+}
+
+object PyArgumentListSemanticChecker : InspectionBasedSemanticChecker(PyArgumentListInspection()) {
+  override fun convertInspectionsResults(originalPsi: PsiFile,
+                                         problemDescriptors: List<ProblemDescriptor>,
+                                         offset: Int,
+                                         prefix: String,
+                                         suggestion: String): List<CorrectnessError> =
+    problemDescriptors.mapNotNull { problemDescriptor ->
+      val location = getLocationInSuggestion(problemDescriptor, offset, prefix, suggestion) ?: return@mapNotNull null
+    if (problemDescriptor.highlightType == ProblemHighlightType.INFORMATION) {
+      return@mapNotNull null
+    }
+    val elementType = problemDescriptor.psiElement.node.elementType
+    if (elementType === PyTokenTypes.RPAR) {
+      return@mapNotNull null
+    }
+      CorrectnessError(location, Severity.CRITICAL)
+  }
+}
+
+object PyRedeclarationSemanticChecker : InspectionBasedSemanticChecker(PyRedeclarationInspection()) {
+  override fun convertInspectionsResults(originalPsi: PsiFile,
+                                         problemDescriptors: List<ProblemDescriptor>,
+                                         offset: Int,
+                                         prefix: String,
+                                         suggestion: String): List<CorrectnessError> =
+    problemDescriptors.mapNotNull { problemDescriptor ->
+      val location = getLocationInSuggestion(problemDescriptor, offset, prefix, suggestion) ?: return@mapNotNull null
+      CorrectnessError(location, Severity.CRITICAL)
+  }
+}
--- a/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/finalizer/PythonSuggestionFinalizer.kt
+++ b/python/src/com/jetbrains/python/codeInsight/mlcompletion/correctness/finalizer/PythonSuggestionFinalizer.kt
@@ -0,0 +1,354 @@
+package com.jetbrains.python.codeInsight.mlcompletion.correctness.finalizer
+
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiWhiteSpace
+import com.intellij.psi.impl.source.tree.LeafPsiElement
+import com.intellij.psi.util.*
+import com.jetbrains.python.PyTokenTypes
+import com.jetbrains.python.PythonLanguage
+import com.jetbrains.python.psi.*
+import com.intellij.platform.ml.impl.correctness.finalizer.SuggestionFinalizerBase
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.ScheduledForRemoval
+@Deprecated("Do not use it")
+class PythonSuggestionFinalizer : SuggestionFinalizerBase(PythonLanguage.INSTANCE) {
+  override fun getTraverseOrder(insertedElement: PsiElement): Sequence<PsiElement> {
+    return sequence {
+      val lastNotComment = generateSequence(insertedElement) { it.prevLeaf() }.firstOrNull { it.text.isNotBlank() && it !is PsiComment }
+                           ?: return@sequence
+      if (insertedElement is PsiComment) {
+        yield(insertedElement)
+      }
+      if (lastNotComment.node.elementType == PyTokenTypes.BACKSLASH) {
+        val lastNotBackslash = generateSequence(lastNotComment) { it.prevLeaf() }.firstOrNull {
+          it.node.elementType != PyTokenTypes.BACKSLASH && it !is PsiWhiteSpace
+        }
+        if (lastNotBackslash != null) {
+          yield(lastNotComment)
+          yieldAll(lastNotBackslash.parents(withSelf = true))
+          return@sequence
+        }
+      }
+      yieldAll(lastNotComment.parents(withSelf = true))
+    }
+  }
+
+  override fun getFinalizationCandidate(element: PsiElement): String? = when (element) {
+    is PyReferenceExpression -> referenceExpression(element)
+    is PyStringElement -> stringElement(element)
+    is PyFStringFragment -> fStringFragment(element)
+    is PyTryPart, is PyFinallyPart -> tryAndFinallyPart(element as PyStatementPart)
+    is PyExceptPart -> exceptPart(element)
+    is PyTryExceptStatement -> tryExceptStatement(element)
+    is PyArgumentList -> argumentList(element)
+    is PyDecorator -> decorator(element)
+    is PyForPart -> forPart(element)
+    is PyWhilePart -> whilePart(element)
+    is PyIfPart, is PyElsePart -> ifElsePart(element)
+    is PyConditionalExpression -> conditionalExpression(element)
+    is PyPrefixExpression -> prefixExpression(element)
+    is PyBinaryExpression -> binaryExpression(element)
+    is PyAssignmentStatement -> assignmentStatement(element)
+    is PyAugAssignmentStatement -> augAssignmentStatement(element)
+    is PyWithStatement -> withStatement(element)
+    is PySubscriptionExpression -> subscriptionExpression(element)
+    is PySliceExpression -> sliceExpression(element)
+    is PyComprehensionElement -> comprehensionElement(element)
+    is PyTupleExpression, is PyParenthesizedExpression -> tupleAndParenthesizedExpression(element)
+    is PyListLiteralExpression -> listLiteral(element)
+    is PySetLiteralExpression -> setLiteral(element)
+    is PyDictLiteralExpression -> dictLiteral(element)
+    is PyImportStatement -> importStatement(element)
+    is PyFromImportStatement -> fromImport(element)
+    is PyClass -> classDeclaration(element)
+    is PyFunction -> functionDeclaration(element)
+    is PyNamedParameter -> namedParameter(element)
+    is PyStarArgument -> starArgument(element)
+    is PyLambdaExpression -> lambdaExpression(element)
+    is PyAssertStatement -> assertStatement(element)
+    is PyDelStatement -> delStatement(element)
+    is PyAnnotation -> annotation(element)
+    is PsiComment -> psiComment()
+    is LeafPsiElement -> leafPsi(element)
+    else -> null
+  }
+
+  private fun referenceExpression(element: PyReferenceExpression): String = when {
+    element.text.endsWith('.') -> "x"
+    else -> ""
+  }
+
+  private fun stringElement(element: PyStringElement): String = when {
+    !element.isTerminated -> " " + element.quote
+    else -> ""
+  }
+
+  private fun fStringFragment(element: PyFStringFragment): String = when {
+    element.expression == null -> "x}"
+    element.closingBrace == null -> "}"
+    else -> ""
+  }
+
+  private fun tryAndFinallyPart(element: PyStatementPart): String {
+    val indent = PyIndentUtil.getElementIndent(element)
+    return when {
+      element.node.findChildByType(PyTokenTypes.COLON) == null -> ": pass\n"
+      element.statementList.statements.isEmpty() -> "$indent pass\n"
+      else -> ""
+    }
+  }
+
+  private fun exceptPart(element: PyExceptPart): String {
+    val indent = PyIndentUtil.getElementIndent(element)
+    return when {
+      element.node.findChildByType(PyTokenTypes.AS_KEYWORD) != null && element.target == null -> " x: pass\n"
+      element.node.findChildByType(PyTokenTypes.COLON) == null -> ": pass\n"
+      element.statementList.statements.isEmpty() -> "$indent pass\n"
+      else -> ""
+    }
+  }
+
+  private fun tryExceptStatement(element: PyTryExceptStatement): String = when {
+    element.exceptParts.isEmpty() && element.finallyPart == null -> {
+      val indent = PyIndentUtil.getElementIndent(element)
+      "\n${indent}except: pass"
+    }
+    else -> ""
+  }
+
+  private fun argumentList(element: PyArgumentList): String {
+    val lastArgument = element.arguments.lastOrNull()
+    return when {
+      lastArgument is PyKeywordArgument && lastArgument.valueExpression == null -> "x)"
+      element.closingParen == null -> ")"
+      else -> ""
+    }
+  }
+
+  private fun decorator(element: PyDecorator): String = when (element.expression) {
+    null -> "x\n"
+    else -> ""
+  }
+
+  private fun forPart(element: PyForPart): String {
+    val indent = PyIndentUtil.getElementIndent(element)
+    return when {
+      element.target == null -> " x in x: pass"
+      element.source == null && element.node.findChildByType(PyTokenTypes.IN_KEYWORD) == null -> " in x: pass"
+      element.source == null -> " x: pass"
+      element.node.findChildByType(PyTokenTypes.COLON) == null -> ": pass"
+      (element.lastChild as? PyStatementList)?.statements?.isEmpty() == true -> "$indent pass"
+      else -> ""
+    }
+  }
+
+  private fun whilePart(element: PyWhilePart): String {
+    val indent = PyIndentUtil.getElementIndent(element)
+    return when {
+      element.condition == null -> " x: pass"
+      element.node.findChildByType(PyTokenTypes.COLON) == null -> ": pass"
+      (element.lastChild as? PyStatementList)?.statements?.isEmpty() == true -> "$indent pass"
+      else -> ""
+    }
+  }
+
+  private fun ifElsePart(element: PsiElement): String {
+    val indent = PyIndentUtil.getElementIndent(element)
+    return when {
+      element is PyIfPart && element.condition == null -> " x: pass"
+      element.node.findChildByType(PyTokenTypes.COLON) == null -> ": pass"
+      (element.lastChild as? PyStatementList)?.statements?.isEmpty() == true -> "$indent pass"
+      else -> ""
+    }
+  }
+
+  private fun conditionalExpression(element: PyConditionalExpression): String = when {
+    element.condition == null -> " x else x"
+    element.falsePart == null && element.node.findChildByType(PyTokenTypes.ELSE_KEYWORD) == null -> " else x"
+    element.falsePart == null -> " x"
+    else -> ""
+  }
+
+  private fun prefixExpression(element: PyPrefixExpression): String = when (element.operand) {
+    null -> "x"
+    else -> ""
+  }
+
+  private fun binaryExpression(element: PyBinaryExpression): String = when (element.rightExpression) {
+    null -> " x"
+    else -> ""
+  }
+
+  private fun assignmentStatement(element: PyAssignmentStatement): String = when (element.assignedValue) {
+    null -> " x"
+    else -> ""
+  }
+
+  private fun augAssignmentStatement(element: PyAugAssignmentStatement): String = when (element.value) {
+    null -> " x"
+    else -> ""
+  }
+
+  private fun withStatement(element: PyWithStatement): String {
+    val indent = PyIndentUtil.getElementIndent(element)
+    return when {
+      element.withItems.isEmpty() -> " x: pass"
+      element.withItems.last().node.findChildByType(
+        PyTokenTypes.AS_KEYWORD) != null && element.withItems.last().target == null -> " x: pass"
+      element.node.findChildByType(PyTokenTypes.COLON) == null -> ": pass"
+      element.statementList.statements.isEmpty() -> "$indent pass"
+      else -> ""
+    }
+  }
+
+  private fun subscriptionExpression(element: PySubscriptionExpression): String = when {
+    element.indexExpression == null -> "x]"
+    element.node.findChildByType(PyTokenTypes.RBRACKET) == null -> "]"
+    else -> ""
+  }
+
+  private fun sliceExpression(element: PySliceExpression): String = when {
+    element.node.findChildByType(PyTokenTypes.RBRACKET) == null -> "]"
+    else -> ""
+  }
+
+  private fun comprehensionElement(element: PyComprehensionElement): String {
+    val (closeBracketType, closeBracketText) = when (element) {
+      is PyListCompExpression -> PyTokenTypes.RBRACKET to "]"
+      is PyGeneratorExpression -> {
+        val brace = if (element.firstChild.node.elementType == PyTokenTypes.LPAR) ")" else ""
+        PyTokenTypes.RPAR to brace
+      }
+      else -> PyTokenTypes.RBRACE to "}"
+    }
+
+    return when {
+      element.forComponents.isEmpty() && element.node.findChildByType(PyTokenTypes.FOR_KEYWORD)!!.psi.siblings(
+        withSelf = false).none { it is PyExpression } -> " x in x$closeBracketText"
+      element.forComponents.isEmpty() && element.node.findChildByType(
+        PyTokenTypes.IN_KEYWORD) == null -> " in x$closeBracketText"
+      element.forComponents.isEmpty() -> " x$closeBracketText"
+      element.ifComponents.isEmpty() && element.node.findChildByType(
+        PyTokenTypes.IF_KEYWORD) != null -> " x$closeBracketText"
+      element.node.findChildByType(closeBracketType) == null -> closeBracketText
+      else -> ""
+    }
+  }
+
+  private fun tupleAndParenthesizedExpression(element: PsiElement): String = when {
+    element.firstChild.node.elementType == PyTokenTypes.LPAR && element.lastChild.node.elementType != PyTokenTypes.RPAR -> ")"
+    else -> ""
+  }
+
+  private fun listLiteral(element: PyListLiteralExpression): String = when {
+    element.lastChild.node.elementType != PyTokenTypes.RBRACKET -> "]"
+    else -> ""
+  }
+
+  private fun setLiteral(element: PySetLiteralExpression): String = when {
+    element.lastChild.node.elementType != PyTokenTypes.RBRACE -> "}"
+    else -> ""
+  }
+
+  private fun dictLiteral(element: PyDictLiteralExpression): String {
+    val nonLeafChildren = element.children.filter { child -> child.firstChild != null }
+    val lastExpr = element.elements.lastOrNull()
+    return when {
+      element.node.findChildByType(PyTokenTypes.COLON) != null || (lastExpr != null && lastExpr.value == null) -> "x}"
+      nonLeafChildren.size > element.elements.size -> ":x}"
+      element.lastChild.node.elementType != PyTokenTypes.RBRACE -> "}"
+      else -> ""
+    }
+  }
+
+  private fun importStatement(element: PyImportStatement): String {
+    val lastImport = element.importElements.lastOrNull()
+    val commaAfterLastImport = lastImport?.siblings(withSelf = false)?.firstOrNull { it.node.elementType == PyTokenTypes.COMMA }
+    return when {
+      lastImport == null -> " x"
+      commaAfterLastImport != null -> " x"
+      lastImport.node.findChildByType(PyTokenTypes.AS_KEYWORD) != null && lastImport.asNameElement == null -> " x"
+      else -> ""
+    }
+  }
+
+  private fun fromImport(element: PyFromImportStatement): String {
+    val lastImport = element.importElements.lastOrNull()
+    return when {
+      element.importSource == null && element.node.findChildByType(PyTokenTypes.IMPORT_KEYWORD) == null -> " x import x"
+      element.node.findChildByType(PyTokenTypes.IMPORT_KEYWORD) == null -> " import x"
+      element.leftParen == null && element.starImportElement == null && lastImport != null && lastImport.importedQName == null -> " x"
+      element.leftParen == null && element.starImportElement == null && lastImport != null && lastImport.node.findChildByType(
+        PyTokenTypes.AS_KEYWORD) != null && lastImport.asNameElement == null -> " x"
+      element.leftParen != null && element.rightParen == null && lastImport != null && lastImport.importedQName == null -> " x)"
+      element.leftParen != null && element.rightParen == null && lastImport != null && lastImport.node.findChildByType(
+        PyTokenTypes.AS_KEYWORD) != null && lastImport.asNameElement == null -> " x)"
+      element.leftParen != null && element.rightParen == null -> ")"
+      else -> ""
+    }
+  }
+
+  private fun classDeclaration(element: PyClass): String {
+    val indent = PyIndentUtil.getElementIndent(element)
+    return when {
+      element.nameNode == null -> " x: pass"
+      element.node.findChildByType(PyTokenTypes.COLON) == null -> ": pass"
+      element.statementList.statements.isEmpty() -> "$indent pass"
+      else -> ""
+    }
+  }
+
+  private fun functionDeclaration(element: PyFunction): String {
+    val indent = PyIndentUtil.getElementIndent(element)
+    return when {
+      element.node.findChildByType(PyTokenTypes.DEF_KEYWORD) == null -> "\n${indent}def x(): pass"
+      element.nameNode == null -> " x(): pass"
+      element.parameterList.node.findChildByType(PyTokenTypes.LPAR) == null -> "(): pass"
+      element.parameterList.node.findChildByType(PyTokenTypes.RPAR) == null -> "): pass"
+      element.node.findChildByType(PyTokenTypes.COLON) == null -> ": pass"
+      element.statementList.statements.isEmpty() -> "$indent pass"
+      else -> ""
+    }
+  }
+
+  private fun namedParameter(element: PyNamedParameter): String = when {
+    element.node.findChildByType(PyTokenTypes.EQ) != null && !element.hasDefaultValue() -> "x"
+    else -> ""
+  }
+
+  private fun starArgument(element: PyStarArgument): String = when {
+    PsiTreeUtil.getChildOfType(element, PyReferenceExpression::class.java) == null -> "x"
+    else -> ""
+  }
+
+  private fun lambdaExpression(element: PyLambdaExpression): String = when {
+    element.node.findChildByType(PyTokenTypes.COLON) == null -> ": 0"
+    element.body == null -> "0"
+    else -> ""
+  }
+
+  private fun assertStatement(element: PyAssertStatement): String = when {
+    element.arguments.isEmpty() -> " x"
+    else -> ""
+  }
+
+  private fun delStatement(element: PyDelStatement): String = when {
+    element.targets.isEmpty() -> " x"
+    else -> ""
+  }
+
+  private fun annotation(element: PyAnnotation): String = when (element.value) {
+    null -> "x"
+    else -> ""
+  }
+
+  private fun psiComment(): String = "\n"
+
+  private fun leafPsi(element: LeafPsiElement): String = when {
+    element.textContains('\\') -> "\n"
+    element.node.elementType == PyTokenTypes.EXP && element.parent?.parent is PyParameterList -> "x"
+    else -> ""
+  }
+}