diff --git a/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogProcessor.kt b/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogProcessor.kt index c9b2b57ae2..bc68d871e6 100644 --- a/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogProcessor.kt +++ b/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogProcessor.kt @@ -318,7 +318,11 @@ fun getDirectoriesForSysPath( if (element != null) { val directory = element.parent if (directory is PsiDirectory) { - importedPaths.add(directory.virtualFile) + // If we have `import a.b.c` we need to add syspath to module `a` only + val additionalLevel = importTarget.importedQName?.componentCount?.dec() ?: 0 + directory.topParent(additionalLevel)?.let { dir -> + importedPaths.add(dir.virtualFile) + } } } } @@ -329,7 +333,11 @@ fun getDirectoriesForSysPath( importTarget.resolveImportSourceCandidates().forEach { val directory = it.parent if (directory is PsiDirectory ) { - importedPaths.add(directory.virtualFile) + // If we have `from a.b.c import d` we need to add syspath to module `a` only + val additionalLevel = importTarget.importSourceQName?.componentCount?.dec() ?: 0 + directory.topParent(additionalLevel)?.let { dir -> + importedPaths.add(dir.virtualFile) + } } } } diff --git a/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/Utils.kt b/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/Utils.kt index c8560beadb..0b73c99e9a 100644 --- a/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/Utils.kt +++ b/utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/Utils.kt @@ -4,6 +4,7 @@ import com.intellij.openapi.project.Project import com.intellij.openapi.roots.ProjectFileIndex import com.intellij.openapi.vfs.VfsUtil import com.intellij.openapi.vfs.VirtualFile +import com.intellij.psi.PsiDirectory import com.intellij.psi.PsiElement import com.jetbrains.python.psi.PyClass import com.jetbrains.python.psi.PyDecorator @@ -55,4 +56,12 @@ fun fineFunction(function: PyFunction): Boolean = fun fineClass(pyClass: PyClass): Boolean = getAncestors(pyClass).dropLast(1).all { it !is PyClass && it !is PyFunction } && - pyClass.methods.any { fineFunction(it) } \ No newline at end of file + pyClass.methods.any { fineFunction(it) } + +fun PsiDirectory.topParent(level: Int): PsiDirectory? { + var directory: PsiDirectory? = this + repeat(level) { + directory = directory?.parent + } + return directory +}