From d23ed85b5dd27252cd7a604b118f8699d808b993 Mon Sep 17 00:00:00 2001
From: Vyacheslav Tamarin <vyacheslav.tamarin@yandex.ru>
Date: Thu, 22 Dec 2022 09:28:16 +0300
Subject: [PATCH] Collect imported modules paths to sys.path

---
 .../language/python/PythonDialogProcessor.kt       | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

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 afa0ad3a08..0f171a78be 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
@@ -271,6 +271,20 @@ fun getDirectoriesForSysPath(
     if (ancestor != null && !sources.contains(ancestor))
         sources.add(ancestor)
 
+    // Collect sys.path directories with imported modules
+    file.importTargets.forEach { importTarget ->
+        importTarget.multiResolve().forEach {
+            val element = it.element
+            if (element != null) {
+                val directory = element.parent
+                if (directory is PsiDirectory) {
+                    sources.add(directory.virtualFile)
+                }
+            }
+
+        }
+    }
+
     var importPath = ancestor?.let { VfsUtil.getParentDir(VfsUtilCore.getRelativeLocation(file.virtualFile, it)) } ?: ""
     if (importPath != "")
         importPath += "."