PsiElement (@NotNull PsiElement insertParent, @Nullable PsiElement anchor, @Nullable PyImportStatementBase newImport, @Nullable ImportPriority priority) { PsiElement feeler = ImportLocationHelper.getInstance().getSearchStartPosition(anchor, insertParent); if (feeler == null) return null; // skip initial comments and whitespace and try to get just below the last import stmt boolean skippedOverStatements = false; boolean skippedOverDoc = false; PsiElement seeker = feeler; final boolean isInjected = InjectedLanguageManager.getInstance(feeler.getProject()).isInjectedFragment(feeler.getContainingFile()); PyImportStatementBase importAbove = null, importBelow = null; do { if (feeler instanceof PyImportStatementBase existingImport && !isInjected) { if (priority != null && newImport != null) { if (shouldInsertBefore(newImport, existingImport, priority)) { importBelow = existingImport; break; } else { importAbove = existingImport; } } seeker = feeler; feeler = feeler.getNextSibling(); skippedOverStatements = true; } else if (PsiTreeUtil.instanceOf(feeler, PsiWhiteSpace.class, PsiComment.class)) { seeker = feeler; feeler = feeler.getNextSibling(); } else if (PsiTreeUtil.instanceOf(feeler, OuterLanguageElement.class)) { if (skippedOverStatements) { break; } feeler = feeler.getNextSibling(); seeker = feeler; } else if (isAssignmentToModuleLevelDunderName(feeler)) { if (priority == ImportPriority.FUTURE) { seeker = feeler; break; } feeler = feeler.getNextSibling(); seeker = feeler; skippedOverStatements = true; } // maybe we arrived at the doc comment stmt; skip over it, too else if (!skippedOverStatements && !skippedOverDoc && insertParent instanceof PyFile) { // this gives the literal; its parent is the expr seeker may have encountered final PsiElement docElem = DocStringUtil.findDocStringExpression((PyElement)insertParent); if (docElem != null && docElem.getParent() == feeler) { feeler = feeler.getNextSibling(); seeker = feeler; // skip over doc even if there's nothing below it skippedOverDoc = true; } else { break; // not a doc comment, stop on it } } else { break; // some other statement, stop } } while (feeler != null); final ImportPriority priorityAbove = importAbove != null ? getImportPriority(importAbove) : null; final ImportPriority priorityBelow = importBelow != null ? getImportPriority(importBelow) : null; if (newImport != null && (priorityAbove == null || priorityAbove.compareTo(priority) < 0)) { newImport.putCopyableUserData(PythonCodeStyleService.IMPORT_GROUP_BEGIN, true); } if (feeler != null) { var anchorComment = getTopmostBoundComment(feeler); if (anchorComment != null) { seeker = anchorComment; } } if (priorityBelow != null) { // actually not necessary because existing import with higher priority (i.e. lower import group) // probably should have IMPORT_GROUP_BEGIN flag already, but we add it anyway just for safety if (priorityBelow.compareTo(priority) > 0) { importBelow.putCopyableUserData(PythonCodeStyleService.IMPORT_GROUP_BEGIN, true); } else if (priorityBelow == priority) { importBelow.putCopyableUserData(PythonCodeStyleService.IMPORT_GROUP_BEGIN, null); } } return seeker; }