Rewrite logic with system paths and imports in python (#1578)

Author: tamarinvs19

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/PythonDialogProcessor.kt

@@ -12,6 +12,7 @@ import com.intellij.openapi.roots.ProjectFileIndex
 import com.intellij.openapi.ui.Messages
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VfsUtilCore
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiDirectory
 import com.intellij.psi.PsiFileFactory
 import com.jetbrains.python.psi.PyClass
@@ -36,7 +37,6 @@ import org.utbot.python.framework.codegen.PythonCgLanguageAssistant
 import org.utbot.python.utils.RequirementsUtils.installRequirements
 import org.utbot.python.utils.RequirementsUtils.requirements
 import org.utbot.python.utils.camelToSnakeCase
-import java.io.File
 import java.nio.file.Path
 import java.nio.file.Paths
 import kotlin.io.path.Path
@@ -285,43 +285,66 @@ fun getPyCodeFromPyFile(file: PyFile, pythonModule: String): PythonCode? {
     return getFromString(content, file.virtualFile.path, pythonModule = pythonModule)
 }
 
+/*
+ * Returns set of sys paths and tested file import path
+ */
 fun getDirectoriesForSysPath(
     srcModule: Module,
     file: PyFile
 ): Pair<Set<String>, String> {
     val sources = ModuleRootManager.getInstance(srcModule).getSourceRoots(false).toMutableList()
     val ancestor = ProjectFileIndex.getInstance(file.project).getContentRootForFile(file.virtualFile)
-    if (ancestor != null && !sources.contains(ancestor))
+    if (ancestor != null)
         sources.add(ancestor)
 
     // Collect sys.path directories with imported modules
+    val importedPaths = emptyList<VirtualFile>().toMutableList()
+
+    // 1. import <module>
     file.importTargets.forEach { importTarget ->
         importTarget.multiResolve().forEach {
             val element = it.element
             if (element != null) {
                 val directory = element.parent
                 if (directory is PsiDirectory) {
-                    if (sources.any { source ->
-                            val sourcePath = source.canonicalPath
-                            if (source.isDirectory && sourcePath != null) {
-                                directory.virtualFile.canonicalPath?.startsWith(sourcePath) ?: false
-                            } else {
-                                false
-                            }
-                        }) {
-                        sources.add(directory.virtualFile)
-                    }
+                    importedPaths.add(directory.virtualFile)
                 }
             }
         }
     }
 
-    var importPath = ancestor?.let { VfsUtil.getParentDir(VfsUtilCore.getRelativeLocation(file.virtualFile, it)) } ?: ""
-    if (importPath != "")
-        importPath += "."
+    // 2. from <module> import ...
+    file.fromImports.forEach { importTarget ->
+        importTarget.resolveImportSourceCandidates().forEach {
+            val directory = it.parent
+            if (directory is PsiDirectory ) {
+                importedPaths.add(directory.virtualFile)
+            }
+        }
+    }
+
+    // Select modules only from this project
+    importedPaths.forEach {
+        if (it.isProjectSubmodule(ancestor)) {
+            sources.add(it)
+        }
+    }
+
+    val fileName = file.name.removeSuffix(".py")
+    val importPath = ancestor?.let {
+        VfsUtil.getParentDir(
+            VfsUtilCore.getRelativeLocation(file.virtualFile, it)
+        )
+    } ?: ""
+    val importStringPath = listOf(
+        importPath.toPath().joinToString("."),
+        fileName
+    )
+        .filterNot { it.isEmpty() }
+        .joinToString(".")
 
     return Pair(
-        sources.map { it.path.replace("\\", "\\\\") }.toSet(),
-        "${importPath}${file.name}".removeSuffix(".py").toPath().joinToString(".").replace("/", File.separator)
+        sources.map { it.path }.toSet(),
+        importStringPath
     )
 }

utbot-intellij-python/src/main/kotlin/org/utbot/intellij/plugin/language/python/Utils.kt

@@ -2,6 +2,7 @@ package org.utbot.intellij.plugin.language.python
 
 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.PsiElement
 import org.utbot.python.utils.RequirementsUtils
@@ -30,6 +31,10 @@ fun generateRandomString(length: Int): String {
         .joinToString("")
 }
 
+fun VirtualFile.isProjectSubmodule(ancestor: VirtualFile?): Boolean {
+    return VfsUtil.isUnder(this, setOf(ancestor).toMutableSet())
+}
+
 fun checkModuleIsInstalled(pythonPath: String, moduleName: String): Boolean {
     return RequirementsUtils.requirementsAreInstalled(pythonPath, listOf(moduleName))
-}
+}

utbot-python/src/main/kotlin/org/utbot/python/PythonEvaluation.kt

@@ -73,8 +73,8 @@ fun startEvaluationProcess(input: EvaluationInput): EvaluationProcess {
         input.directoriesForSysPath,
         input.moduleToImport,
         input.additionalModulesToImport,
-        fileForOutput.path.replace("\\", "\\\\"),
-        coverageDatabasePath.absolutePath.replace("\\", "\\\\")
+        fileForOutput.path,
+        coverageDatabasePath.path,
     )
     val fileWithCode = TemporaryFileManager.createTemporaryFile(
         runCode,

utbot-python/src/main/kotlin/org/utbot/python/framework/codegen/model/PythonCodeGenerator.kt

@@ -33,6 +33,7 @@ import org.utbot.python.framework.codegen.model.constructor.visitor.CgPythonRend
 import org.utbot.python.framework.codegen.model.tree.CgPythonDict
 import org.utbot.python.framework.codegen.model.tree.CgPythonFunctionCall
 import org.utbot.python.framework.codegen.model.tree.CgPythonList
+import org.utbot.python.framework.codegen.toPythonRawString
 
 class PythonCodeGenerator(
     classUnderTest: ClassId,
@@ -128,9 +129,9 @@ class PythonCodeGenerator(
             arguments.associateBy { argument -> CgLiteral(pythonStrClassId, "'${argument.name}'") }
         )
 
-        val fullpath = CgLiteral(pythonStrClassId, "'${method.moduleFilename.replace("\\", "\\\\")}'")
-        val outputPath = CgLiteral(pythonStrClassId, "'$fileForOutputName'")
-        val databasePath = CgLiteral(pythonStrClassId, "'$coverageDatabasePath'")
+        val fullpath = CgLiteral(pythonStrClassId, method.moduleFilename.toPythonRawString())
+        val outputPath = CgLiteral(pythonStrClassId, fileForOutputName.toPythonRawString())
+        val databasePath = CgLiteral(pythonStrClassId, coverageDatabasePath.toPythonRawString())
 
         val executorCall = CgPythonFunctionCall(
             pythonNoneClassId,

utbot-python/src/main/kotlin/org/utbot/python/framework/codegen/model/constructor/visitor/CgPythonRenderer.kt

@@ -63,6 +63,7 @@ import org.utbot.python.framework.api.python.PythonClassId
 import org.utbot.python.framework.api.python.pythonBuiltinsModuleName
 import org.utbot.python.framework.api.python.util.pythonAnyClassId
 import org.utbot.python.framework.codegen.model.tree.*
+import org.utbot.python.framework.codegen.toPythonRawString
 
 internal class CgPythonRenderer(
     context: CgRendererContext,
@@ -291,7 +292,7 @@ internal class CgPythonRenderer(
 
     fun renderPythonImport(pythonImport: PythonImport) {
         if (pythonImport is PythonSysPathImport) {
-            println("sys.path.append('${pythonImport.sysPath}')")
+            println("sys.path.append(${pythonImport.sysPath.toPythonRawString()})")
         } else if (pythonImport.moduleName == null) {
             println("import ${pythonImport.importName}")
         } else {
@@ -541,4 +542,5 @@ internal class CgPythonRenderer(
             .replace("'", "\\'")
             .replace("\\f", "\\u000C")
             .replace("\\xxx", "\\\u0058\u0058\u0058")
-}
+}
+

utbot-python/src/main/kotlin/org/utbot/python/framework/codegen/utils/StringUtils.kt

@@ -0,0 +1,6 @@
+package org.utbot.python.framework.codegen
+
+
+fun String.toPythonRawString(): String {
+    return "r'${this}'"
+}