From 14203deb33b38081a428f7b7babe59ee14d25cf8 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Tue, 24 Dec 2019 19:49:20 +0300
Subject: [PATCH] PY-23968 Simplify the implementation using
 PyFromImportStatement#addBefore

Previously not functioning because of the bugs in the implementation of
PyFromImportStatementImpl#addInternal.

GitOrigin-RevId: fb70b6472f50ad98e5ba0cadda453b081ca0d9f1
---
 .../codeInsight/imports/AddImportHelper.java       | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java
index 62f0326525f4..dd2ad5866ea6 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java
@@ -588,24 +588,20 @@ public class AddImportHelper {
     final PythonCodeStyleService pyCodeStyle = PythonCodeStyleService.getInstance();
     final boolean shouldSort = pyCodeStyle.isOptimizeImportsSortImports(file) && pyCodeStyle.isOptimizeImportsSortNamesInFromImports(file);
 
-    PyImportElement precedingNameElement = null;
+    PyImportElement followingNameElement = null;
     for (PyImportElement existingNameElement : fromImport.getImportElements()) {
       final String existingName = Objects.toString(existingNameElement.getImportedQName(), "");
       if (name.equals(existingName) && Comparing.equal(asName, existingNameElement.getAsName())) {
         return false;
       }
-      if (shouldSort && nameComparator.compare(existingName, name) < 0) {
-        precedingNameElement = existingNameElement;
+      if (shouldSort && followingNameElement == null && nameComparator.compare(existingName, name) > 0) {
+        followingNameElement = existingNameElement;
       }
     }
     final PyElementGenerator generator = PyElementGenerator.getInstance(fromImport.getProject());
     final PyImportElement newNameElement = generator.createImportElement(LanguageLevel.forElement(fromImport), name, asName);
-    if (shouldSort) {
-      fromImport.addAfter(newNameElement, precedingNameElement);
-    }
-    else {
-      fromImport.add(newNameElement);
-    }
+    // addBefore(newNameElement, null) is the same as inserting at the end
+    fromImport.addBefore(newNameElement, followingNameElement);
     // May need to add parentheses, trailing comma, etc.
     CodeStyleManager.getInstance(fromImport.getProject()).reformat(fromImport);
     return true;