relative imports should be inserted instead of absolute ones when possible (PY-6054)

If the file already has relative imports, next one inserted via completion/auto-import has a chance to be relative. For that new import has to have the number of dots <= than the max of the file's relative import and a number provided by a registry value in 'python.relative.import.depth', that is 3 by default. It also has to be located in the same root package, since python doesn't allow relative imports to step outside of it. The file must not have "if __name__ == '__main__'" in it, since it suggests that it is intended to be runnable and relative imports will not work.
Relative import statements can also be updated with new import elements, if the user's codestyle allows it.

GitOrigin-RevId: 766a03e2252bb1f8193a2156831cc8fdf0a7aced

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

@@ -9,6 +9,7 @@ import com.intellij.openapi.projectRoots.Sdk;
 import com.intellij.openapi.roots.ProjectFileIndex;
 import com.intellij.openapi.roots.ProjectRootManager;
 import com.intellij.openapi.util.Comparing;
+import com.intellij.openapi.util.registry.Registry;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.*;
@@ -29,6 +30,7 @@ import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
 import com.jetbrains.python.pyi.PyiFile;
 import com.jetbrains.python.pyi.PyiUtil;
 import com.jetbrains.python.sdk.PythonSdkUtil;
+import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -514,14 +516,35 @@ public class AddImportHelper {
                                                        @Nullable String asName,
                                                        @Nullable ImportPriority priority,
                                                        @Nullable PsiElement anchor) {
-    final List<PyFromImportStatement> existingImports = ((PyFile)file).getFromImports();
+    final PyFile pyFile = (PyFile)file;
+    final List<PyFromImportStatement> existingImports = pyFile.getFromImports();
+
+    int relativeLevel = 0;
+    PyRelativeImportData importData = null;
+    if (priority == ImportPriority.PROJECT) {
+      importData = PyRelativeImportData.fromString(from, pyFile);
+      if (importData != null) {
+        relativeLevel = importData.getRelativeLevel();
+      }
+    }
+
     if (!PythonCodeStyleService.getInstance().isOptimizeImportsAlwaysSplitFromImports(file)) {
       for (PyFromImportStatement existingImport : existingImports) {
         if (existingImport.isStarImport()) {
           continue;
         }
         final String existingSource = Objects.toString(existingImport.getImportSourceQName(), "");
-        if (from.equals(existingSource) && existingImport.getRelativeLevel() == 0) {
+
+        boolean updateExisting = false;
+        final int currentRelativeLevel = existingImport.getRelativeLevel();
+        if (currentRelativeLevel == 0) {
+          updateExisting = from.equals(existingSource);
+        }
+        else if (relativeLevel != 0) {
+          updateExisting = currentRelativeLevel == relativeLevel && existingSource.equals(importData.getRelativeLocation());
+        }
+
+        if (updateExisting) {
           for (PyImportElement el : existingImport.getImportElements()) {
             final String existingName = Objects.toString(el.getImportedQName(), "");
             if (name.equals(existingName) && Comparing.equal(asName, el.getAsName())) {
@@ -537,6 +560,21 @@ public class AddImportHelper {
         }
       }
     }
+
+    if (!PyUtil.hasIfNameEqualsMain(pyFile)) {
+      final int maxRelativeLevelInFile = StreamEx.of(pyFile.getFromImports())
+        .mapToInt(PyFromImportStatement::getRelativeLevel)
+        .max()
+        .orElse(0);
+
+      if (maxRelativeLevelInFile > 0 && relativeLevel > 0) {
+        int maxAllowedDepth = Math.max(maxRelativeLevelInFile, Registry.intValue("python.relative.import.depth"));
+        if (maxAllowedDepth >= relativeLevel) {
+          from = importData.getLocationWithDots();
+        }
+      }
+    }
+
     addFromImportStatement(file, from, name, asName, priority, anchor);
     return true;
   }

python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/PyRelativeImportData.kt

@@ -0,0 +1,36 @@
+package com.jetbrains.python.codeInsight.imports
+
+import com.intellij.psi.util.QualifiedName
+import com.jetbrains.python.PyNames
+import com.jetbrains.python.psi.PyFile
+import com.jetbrains.python.psi.resolve.QualifiedNameFinder
+
+/**
+ * @author Aleksei Kniazev
+ */
+class PyRelativeImportData private constructor(val relativeLocation: String, val relativeLevel: Int)  {
+
+  val locationWithDots: String
+    get() = ".".repeat(relativeLevel) + relativeLocation
+
+  companion object {
+    @JvmStatic
+    fun fromString(location: String, file: PyFile): PyRelativeImportData? {
+      val qName = QualifiedName.fromDottedString(location)
+      val fileQName = QualifiedNameFinder.findCanonicalImportPath(file, null) ?: return null
+
+      if (qName.firstComponent != fileQName.firstComponent) return null
+
+      val common = fileQName.components.asSequence()
+        .zip(qName.components.asSequence()) { s1, s2 -> s1 == s2 }
+        .takeWhile { it }
+        .count()
+
+      val remainingFileQname = fileQName.removeHead(common)
+      val remainingQname = qName.removeHead(common)
+      val implicitLevel = if (file.name == PyNames.INIT_DOT_PY) 1 else 0
+
+      return PyRelativeImportData(remainingQname.toString(), remainingFileQname.componentCount + implicitLevel)
+    }
+  }
+}

python/python-psi-impl/src/com/jetbrains/python/psi/PyUtil.java

@@ -290,6 +290,14 @@ public class PyUtil {
     return element instanceof PyTargetExpression && ScopeUtil.getScopeOwner(element) instanceof PyClass;
   }
 
+  public static boolean hasIfNameEqualsMain(@NotNull PyFile file) {
+    final PyIfStatement dunderMain = SyntaxTraverser.psiApi()
+      .children(file)
+      .filterMap(psi -> psi instanceof PyIfStatement ? ((PyIfStatement)psi) : null)
+      .find(ifStatement -> isIfNameEqualsMain(ifStatement));
+    return dunderMain != null;
+  }
+
   public static boolean isIfNameEqualsMain(PyIfStatement ifStatement) {
     final PyExpression condition = ifStatement.getIfPart().getCondition();
     return isNameEqualsMain(condition);

python/src/META-INF/python-core-common.xml

@@ -619,6 +619,7 @@
     <applicationService serviceInterface="com.jetbrains.python.PythonRuntimeService" serviceImplementation="com.jetbrains.python.PythonRuntimeServiceImpl" overrides="true"/>
     <applicationService serviceInterface="com.jetbrains.python.PythonCodeStyleService" serviceImplementation="com.jetbrains.python.PythonCodeStyleServiceImpl" overrides="true"/>
 
+    <registryKey key="python.relative.import.depth" defaultValue="3" description="Specifies default acceptable number of dots in a relative import statement"/>
   </extensions>
 
   <extensionPoints>

python/testData/addImport/relativeImportFromAlreadyImportedModule/foo/bar.py

@@ -0,0 +1,6 @@
+def oh():
+    pass
+
+
+def no():
+    pass

python/testData/addImport/relativeImportFromAlreadyImportedModule/foo/test.after.py

@@ -0,0 +1 @@
+from .bar import oh, no

python/testData/addImport/relativeImportFromAlreadyImportedModule/foo/test.py

@@ -0,0 +1 @@
+from .bar import oh

python/testData/addImport/relativeImportFromSamePackage/foo/bar.py

@@ -0,0 +1,2 @@
+def bar_func():
+    pass

python/testData/addImport/relativeImportFromSamePackage/foo/baz.py

@@ -0,0 +1,2 @@
+def baz_func():
+    pass

python/testData/addImport/relativeImportFromSamePackage/foo/test.after.py

@@ -0,0 +1,2 @@
+from .baz import baz_func
+from .bar import bar_func

python/testData/addImport/relativeImportFromSamePackage/foo/test.py

@@ -0,0 +1 @@
+from .bar import bar_func

python/testData/addImport/relativeImportInFIleWithMain/foo/bar.py

@@ -0,0 +1,2 @@
+def bar_func():
+    pass

python/testData/addImport/relativeImportInFIleWithMain/foo/baz.py

@@ -0,0 +1,2 @@
+def baz_func():
+    pass

python/testData/addImport/relativeImportInFIleWithMain/foo/test.after.py

@@ -0,0 +1,5 @@
+from foo.baz import baz_func
+from .bar import bar_func
+
+if __name__ == '__main__':
+    pass

python/testData/addImport/relativeImportInFIleWithMain/foo/test.py

@@ -0,0 +1,4 @@
+from .bar import bar_func
+
+if __name__ == '__main__':
+    pass

python/testData/addImport/relativeImportInInitFile/foo/bar.py

@@ -0,0 +1,2 @@
+def bar_func():
+    pass

python/testData/addImport/relativeImportInInitFile/foo/src/baz.py

@@ -0,0 +1,2 @@
+def func():
+    pass

python/testData/addImport/relativeImportInInitFile/foo/test/__init__.after.py

@@ -0,0 +1,2 @@
+from ..src.baz import func
+from ..bar import bar_func

python/testData/addImport/relativeImportInInitFile/foo/test/__init__.py

@@ -0,0 +1 @@
+from ..bar import bar_func

python/testData/addImport/relativeImportTooDeep/pkg1/foo.py

@@ -0,0 +1,2 @@
+def foo_func():
+    pass

python/testData/addImport/relativeImportTooDeep/pkg1/pkg2/pkg3/__init__.py

@@ -0,0 +1,2 @@
+def bar_func():
+    pass

python/testData/addImport/relativeImportTooDeep/pkg1/pkg2/pkg3/pkg4/test.after.py

@@ -0,0 +1,2 @@
+from pkg1.foo import foo_func
+from .. import bar_func

python/testData/addImport/relativeImportTooDeep/pkg1/pkg2/pkg3/pkg4/test.py

@@ -0,0 +1 @@
+from .. import bar_func

python/testData/addImport/relativeImportTooDeepWithSameLevelUsed/pkg1/bar.py

@@ -0,0 +1,2 @@
+def bar_func():
+    pass

python/testData/addImport/relativeImportTooDeepWithSameLevelUsed/pkg1/foo.py

@@ -0,0 +1,2 @@
+def foo_func():
+    pass

python/testData/addImport/relativeImportTooDeepWithSameLevelUsed/pkg1/pkg2/pkg3/pkg4/test.after.py

@@ -0,0 +1,2 @@
+from ....foo import foo_func
+from ....bar import bar_func

python/testData/addImport/relativeImportTooDeepWithSameLevelUsed/pkg1/pkg2/pkg3/pkg4/test.py

@@ -0,0 +1 @@
+from ....bar import bar_func

python/testData/addImport/relativeImportWithDotsOnly/foo/bar/baz.py

@@ -0,0 +1,2 @@
+def baz_func():
+    pass

python/testData/addImport/relativeImportWithDotsOnly/foo/bar/test.after.py

@@ -0,0 +1,2 @@
+from .. import lib
+from .baz import baz_func

python/testData/addImport/relativeImportWithDotsOnly/foo/bar/test.py

@@ -0,0 +1 @@
+from .baz import baz_func

python/testSrc/com/jetbrains/python/PyAddImportTest.java

@@ -132,6 +132,41 @@ public class PyAddImportTest extends PyTestCase {
     doAddFromImport("collections", "OrderedDict", BUILTIN);
   }
 
+  // PY-6054
+  public void testRelativeImportFromSamePackage() {
+    doTestRelativeImport("foo.baz", "baz_func", "foo/test");
+  }
+
+  // PY-6054
+  public void testRelativeImportFromAlreadyImportedModule() {
+    doTestRelativeImport("foo.bar", "no", "foo/test");
+  }
+
+  // PY-6054
+  public void testRelativeImportInInitFile() {
+    doTestRelativeImport("foo.src.baz", "func", "foo/test/__init__");
+  }
+
+  // PY-6054
+  public void testRelativeImportInFileWithMain() {
+    doTestRelativeImport("foo.baz", "baz_func", "foo/test");
+  }
+
+  // PY-6054
+  public void testRelativeImportTooDeep() {
+    doTestRelativeImport("pkg1.foo", "foo_func", "pkg1/pkg2/pkg3/pkg4/test");
+  }
+
+  // PY-6054
+  public void testRelativeImportTooDeepWithSameLevelUsed() {
+    doTestRelativeImport("pkg1.foo", "foo_func", "pkg1/pkg2/pkg3/pkg4/test");
+  }
+
+  // PY-6054
+  public void testRelativeImportWithDotsOnly() {
+    doTestRelativeImport("foo", "lib", "foo/bar/test");
+  }
+
   private void doAddOrUpdateFromImport(final String path, final String name, final ImportPriority priority) {
     myFixture.configureByFile(getTestName(true) + ".py");
     WriteCommandAction.runWriteCommandAction(myFixture.getProject(), () -> {
@@ -184,6 +219,16 @@ public class PyAddImportTest extends PyTestCase {
     myFixture.checkResultByFile(getTestName(true) + ".after.py");
   }
 
+  private void doTestRelativeImport(final @NotNull String from, final @NotNull String name, final @NotNull String file) {
+    final String testName = getTestName(true);
+    myFixture.copyDirectoryToProject(testName, "");
+    myFixture.configureByFile(file + ".py");
+    WriteCommandAction.runWriteCommandAction(myFixture.getProject(), () -> {
+      AddImportHelper.addOrUpdateFromImportStatement(myFixture.getFile(), from, name, null, PROJECT, null);
+    });
+    myFixture.checkResultByFile(testName + "/" + file + ".after.py");
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/addImport";