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 }