PY-24044 Evaluate expression and conditional breakpoints should allow adding imports for unresolved names

Added new PyCodeFragmentWithHiddenImports that allows adding imports for unresolved names into a hidden file

GitOrigin-RevId: fb3fa4e491164bdd376d3b25204202034994afa5

python/python-ast/src/com/jetbrains/python/ast/PyAstExpressionCodeFragment.java

@@ -16,8 +16,29 @@
 package com.jetbrains.python.ast;
 
 
+import com.intellij.psi.PsiCodeFragment;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.search.GlobalSearchScope;
 import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
 
 @ApiStatus.Experimental
-public interface PyAstExpressionCodeFragment extends PyAstFile {
+public interface PyAstExpressionCodeFragment extends PyAstFile, PsiCodeFragment {
+  /**
+   * Retrieves the real context of fragment, e.g., if fragment is breakpoint condition,
+   * returns the file in which breakpoint is set.
+   * On the `getContext()` may return a hidden file with imports (and real context is context of that file)
+   *
+   * @return the real context of the element, or null if there is no context
+   */
+  @Nullable
+  default PsiElement getRealContext() {
+    return getContext();
+  }
+
+  @Override
+  default void forceResolveScope(GlobalSearchScope scope) {}
+
+  @Override
+  default GlobalSearchScope getForcedResolveScope() { return null; }
 }

python/python-psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -298,8 +298,8 @@ public final class PyPsiUtils {
   public static PsiElement getRealContext(@NotNull final PsiElement element) {
     assertValid(element);
     final PsiFile file = element.getContainingFile();
-    if (file instanceof PyExpressionCodeFragment) {
-      final PsiElement context = file.getContext();
+    if (file instanceof PyExpressionCodeFragment fragment) {
+      final PsiElement context = fragment.getRealContext();
       return context != null ? context : element;
     }
     else {

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java

@@ -26,6 +26,7 @@ import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
 import com.jetbrains.python.documentation.doctest.PyDocstringFile;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyCodeFragmentWithHiddenImports;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
 import com.jetbrains.python.pyi.PyiFile;
@@ -444,6 +445,9 @@ public final class AddImportHelper {
                                            @Nullable ImportPriority priority,
                                            @Nullable PsiElement anchor,
                                            final @Nullable PsiElement insertBefore) {
+    if (file instanceof PyCodeFragmentWithHiddenImports fragment) {
+      file = fragment.getImportContext();
+    }
     if (!(file instanceof PyFile)) {
       return false;
     }
@@ -794,6 +798,9 @@ public final class AddImportHelper {
    * @see #addOrUpdateFromImportStatement
    */
   public static void addImport(@NotNull PsiNamedElement target, @NotNull PsiFile file, @NotNull PyElement element) {
+    if (file instanceof PyCodeFragmentWithHiddenImports fragment) {
+      file = fragment.getImportContext();
+    }
     if (target instanceof PsiFileSystemItem) {
       addFileSystemItemImport((PsiFileSystemItem)target, file, element);
       return;
@@ -810,6 +817,7 @@ public final class AddImportHelper {
 
     final String path = importPath.toString();
     final ImportPriority priority = getImportPriority(file, toImport);
+
     if (!PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT) {
       addImportStatement(file, path, null, priority, element);
 

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java

@@ -5,13 +5,11 @@ import com.intellij.codeInsight.hint.QuestionAction;
 import com.intellij.lang.injection.InjectedLanguageManager;
 import com.intellij.openapi.command.WriteCommandAction;
 import com.intellij.openapi.project.Project;
-import com.intellij.psi.PsiDocumentManager;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiFileSystemItem;
+import com.intellij.psi.*;
 import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyPsiBundle;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyCodeFragmentWithHiddenImports;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -134,6 +132,9 @@ public class ImportFromExistingAction implements QuestionAction {
     if (manager.isInjectedFragment(file)) {
       file = manager.getTopLevelFile(myTarget);
     }
+    if (file instanceof PyCodeFragmentWithHiddenImports fragment) {
+      file = fragment.getImportContext();
+    }
     // A root-level module or package cannot be imported with a "from" import.
     if (PyUtil.isRoot(item.getFile())) {
       if (myImportLocally) {

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyCodeFragmentWithHiddenImports.kt

@@ -0,0 +1,54 @@
+package com.jetbrains.python.psi.impl
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.ResolveState
+import com.intellij.psi.scope.PsiScopeProcessor
+
+class PyCodeFragmentWithHiddenImports @JvmOverloads constructor(
+  project: Project,
+  name: String,
+  text: CharSequence,
+  isPhysical: Boolean,
+  supportsHiddenImports: Boolean = false,
+) : PyExpressionCodeFragmentImpl(project, name, text, isPhysical) {
+
+  private var myPseudoImportsFragment =
+    if (supportsHiddenImports)
+      PyCodeFragmentWithHiddenImports(project, "imports.py", "", isPhysical, supportsHiddenImports = false).also {
+        it.context = super.getContext()
+        super.setContext(it)
+      }
+    else null
+
+  /**
+   * @return the file where imports should be placed. Either `this` or hidden file for imports
+   */
+  fun getImportContext(): PyCodeFragmentWithHiddenImports = myPseudoImportsFragment ?: this
+  override fun getRealContext(): PsiElement? = getImportContext().context
+
+  override fun setContext(context: PsiElement?) {
+    if (myPseudoImportsFragment != null) {
+      myPseudoImportsFragment!!.context = context
+    } else {
+      super.setContext(context)
+    }
+  }
+
+  override fun processDeclarations(
+    processor: PsiScopeProcessor,
+    state: ResolveState,
+    lastParent: PsiElement?,
+    place: PsiElement,
+  ): Boolean {
+    myPseudoImportsFragment?.processDeclarations(processor, state, lastParent, place)
+    return super.processDeclarations(processor, state, lastParent, place)
+  }
+
+  override fun clone(): PyCodeFragmentWithHiddenImports {
+    val clone = super.clone() as PyCodeFragmentWithHiddenImports
+    clone.myPseudoImportsFragment = myPseudoImportsFragment?.clone()
+    clone.context = clone.myPseudoImportsFragment ?: clone.context
+    return clone
+  }
+}

python/src/com/jetbrains/python/debugger/PyDebuggerEditorsProvider.java

@@ -15,6 +15,7 @@ import com.intellij.xdebugger.evaluation.InlineDebuggerHelper;
 import com.intellij.xdebugger.evaluation.XDebuggerEditorsProvider;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.psi.impl.PyExpressionCodeFragmentImpl;
+import com.jetbrains.python.psi.impl.PyCodeFragmentWithHiddenImports;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -32,7 +33,7 @@ public class PyDebuggerEditorsProvider extends XDebuggerEditorsProvider {
                                           final @Nullable XSourcePosition sourcePosition,
                                           @NotNull EvaluationMode mode) {
     String text = expression.getExpression().trim();
-    final PyExpressionCodeFragmentImpl fragment = new PyExpressionCodeFragmentImpl(project, "fragment.py", text, true);
+    final PyExpressionCodeFragmentImpl fragment = new PyCodeFragmentWithHiddenImports(project, "fragment.py", text, true, true);
 
     // Bind to context
     final PsiElement element = getContextElement(project, sourcePosition);