SSR: optimization in reading configurations
[idea/community.git] / python / src / com / jetbrains / python / testing / PythonUnitTestUtil.java
1 // Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
2 package com.jetbrains.python.testing;
3
4 import com.intellij.execution.Location;
5 import com.intellij.execution.PsiLocation;
6 import com.intellij.openapi.progress.ProgressManager;
7 import com.intellij.openapi.project.Project;
8 import com.intellij.openapi.util.io.FileUtilRt;
9 import com.intellij.openapi.vfs.LocalFileSystem;
10 import com.intellij.openapi.vfs.VirtualFile;
11 import com.intellij.psi.PsiFile;
12 import com.intellij.psi.PsiManager;
13 import com.jetbrains.python.psi.PyClass;
14 import com.jetbrains.python.psi.PyFunction;
15 import com.jetbrains.python.psi.stubs.PyClassNameIndex;
16 import com.jetbrains.python.psi.stubs.PyFunctionNameIndex;
17 import org.jetbrains.annotations.NotNull;
18 import org.jetbrains.annotations.Nullable;
19
20 import java.util.ArrayList;
21 import java.util.List;
22
23 public final class PythonUnitTestUtil {
24   private PythonUnitTestUtil() {}
25
26   public static List<Location> findLocations(@NotNull final Project project,
27                                              @NotNull String fileName,
28                                              @Nullable String className,
29                                              @Nullable String methodName) {
30     if (fileName.contains("%")) {
31       fileName = fileName.substring(0, fileName.lastIndexOf("%"));
32     }
33     final List<Location> locations = new ArrayList<>();
34     if (methodName == null && className == null) {
35       final VirtualFile virtualFile = LocalFileSystem.getInstance().findFileByPath(fileName);
36       if (virtualFile == null) return locations;
37       final PsiFile psiFile = PsiManager.getInstance(project).findFile(virtualFile);
38       if (psiFile != null) {
39         locations.add(new PsiLocation<>(project, psiFile));
40       }
41     }
42
43     if (className != null) {
44       for (PyClass cls : PyClassNameIndex.find(className, project, false)) {
45         ProgressManager.checkCanceled();
46
47         final PsiFile containingFile = cls.getContainingFile();
48         final VirtualFile virtualFile = containingFile.getVirtualFile();
49         final String clsFileName = virtualFile == null ? containingFile.getName() : virtualFile.getPath();
50         final String clsFileNameWithoutExt = FileUtilRt.getNameWithoutExtension(clsFileName);
51         if (!clsFileNameWithoutExt.endsWith(fileName) && !fileName.equals(clsFileName)) {
52           continue;
53         }
54         if (methodName == null) {
55           locations.add(new PsiLocation<>(project, cls));
56         }
57         else {
58           final PyFunction method = cls.findMethodByName(methodName, true, null);
59           if (method == null) {
60             continue;
61           }
62
63           locations.add(new PyPsiLocationWithFixedClass(project, method, cls));
64         }
65       }
66     }
67     else if (methodName != null) {
68       for (PyFunction function : PyFunctionNameIndex.find(methodName, project)) {
69         ProgressManager.checkCanceled();
70         if (function.getContainingClass() == null) {
71           final PsiFile containingFile = function.getContainingFile();
72           final VirtualFile virtualFile = containingFile.getVirtualFile();
73           final String clsFileName = virtualFile == null ? containingFile.getName() : virtualFile.getPath();
74           final String clsFileNameWithoutExt = FileUtilRt.getNameWithoutExtension(clsFileName);
75           if (!clsFileNameWithoutExt.endsWith(fileName)) {
76             continue;
77           }
78           locations.add(new PsiLocation<>(project, function));
79         }
80       }
81     }
82     return locations;
83   }
84 }