com.jetbrains.python.psi.PyFile Java Examples

@Nullable
private static TestLocation testLocation(PsiElement element) {
  PsiFile file = element.getContainingFile();
  if (!(file instanceof PyFile) || !PyTestUtils.isTestFile((PyFile) file)) {
    return null;
  }
  PyClass pyClass = PsiTreeUtil.getParentOfType(element, PyClass.class, false);
  if (pyClass == null || !PyTestUtils.isTestClass(pyClass)) {
    return new TestLocation((PyFile) file, null, null);
  }
  PyFunction pyFunction = PsiTreeUtil.getParentOfType(element, PyFunction.class, false);
  if (pyFunction != null && PyTestUtils.isTestFunction(pyFunction)) {
    return new TestLocation((PyFile) file, pyClass, pyFunction);
  }
  return new TestLocation((PyFile) file, pyClass, null);
}
 

@Nullable
@Override
public BinaryRunContext getRunContext(ConfigurationContext context) {
  Location<?> location = context.getLocation();
  if (location == null) {
    return null;
  }
  PsiElement element = location.getPsiElement();
  PsiFile file = element.getContainingFile();
  if (!(file instanceof PyFile)) {
    return null;
  }
  TargetInfo binaryTarget = getTargetLabel(file);
  if (binaryTarget == null) {
    return null;
  }
  return BinaryRunContext.create(file, binaryTarget);
}
 

private boolean isOrContainsPyTests(PsiElement element) {
  if (new PyTestFinder().isTest(element)) {
    return true;
  }

  if (element instanceof PyFile) {
    PyFile pyFile = (PyFile) element;

    for (PyClass pyClass : pyFile.getTopLevelClasses()) {
      if (PythonUnitTestUtil.isTestClass(pyClass, ThreeState.YES, null)) {
        return true;
      }
    }
  }
  else if (element instanceof PsiDirectory) {
    PsiDirectory directory = (PsiDirectory) element;
    for (PsiFile file : directory.getFiles()) {
      if (isOrContainsPyTests(file)) {
        return true;
      }
    }
  }

  return false;
}
 

public static List<PsiElement> resolveTargetByName(@Nls String name, @NotNull Project project, GlobalSearchScope scope) {
  final PsiManager psiManager = PsiManager.getInstance(project);
  final ArrayList<PsiElement> result = new ArrayList<>();
  final Collection<VirtualFile> containingFiles = FileBasedIndex.getInstance().getContainingFiles(NAME, name, scope);
  for (VirtualFile virtualFile : containingFiles) {
    final PsiFile psiFile = psiManager.findFile(virtualFile);
    if (psiFile instanceof PyFile) {
      final PyReferenceExpression referenceExpression = PantsPsiUtil.findTargetDefinitions((PyFile)psiFile).get(name);
      final PsiPolyVariantReference reference = referenceExpression != null ? referenceExpression.getReference() : null;
      final PsiElement definition = reference != null ? reference.resolve() : null;
      if (definition != null) {
        result.add(definition);
      } else if (referenceExpression != null) {
        // at least something
        result.add(referenceExpression);
      }
    }
  }
  return result;
}
 

@Override
@NotNull
public Map<String, Void> map(@NotNull final FileContent inputData) {
  final PsiFile psiFile = inputData.getPsiFile();
  if (psiFile instanceof PyFile) {
    final Map<String, PyReferenceExpression> targetDefinitions = PantsPsiUtil.findTargetDefinitions((PyFile)psiFile);
    return ContainerUtil.newMapFromKeys(
      targetDefinitions.keySet().iterator(),
      new Convertor<String, Void>() {
        @Override
        public Void convert(String o) {
          return null;
        }
      }
    );
  }
  return Collections.emptyMap();
}
 

@NotNull
@Override
public Collection<KeyUsagePsiElement> getUsages(PsiFile psiFile) {
    if(psiFile instanceof PyFile) {
        PythonEnvironmentCallsVisitor visitor = new PythonEnvironmentCallsVisitor();
        psiFile.acceptChildren(visitor);

        return visitor.getCollectedItems();
    }

    return Collections.emptyList();
}
 

@Override
public List<RatedResolveResult> resolveName(
    PyQualifiedExpression element, TypeEvalContext context) {
  if (!Blaze.isBlazeProject(element.getProject())) {
    return ImmutableList.of();
  }
  PyBuiltinCache cache = getBuiltInCache(element);
  if (cache == null) {
    return ImmutableList.of();
  }
  String referencedName = element.getReferencedName();
  if (referencedName == null) {
    return ImmutableList.of();
  }
  List<RatedResolveResult> result = Lists.newArrayList();
  PyFile bfile = cache.getBuiltinsFile();
  if (bfile != null && !PyUtil.isClassPrivateName(referencedName)) {
    PsiElement resultElement = bfile.getElementNamed(referencedName);
    if (resultElement == null && "__builtins__".equals(referencedName)) {
      resultElement = bfile; // resolve __builtins__ reference
    }
    if (resultElement != null) {
      result.add(
          new ImportedResolveResult(
              resultElement, PyReferenceImpl.getRate(resultElement, context), null));
    }
  }
  return result;
}
 

@Nullable
@Override
public PsiFileAndName toPsiFileAndName(BlazeProjectData projectData, ProjectViewNode<?> node) {
  if (!(node instanceof PsiFileNode)) {
    return null;
  }
  PsiFile psiFile = ((PsiFileNode) node).getValue();
  if (!(psiFile instanceof PyFile)) {
    return null;
  }
  return new PsiFileAndName(psiFile, psiFile.getName());
}
 

@Test
public void testResolveWorkspaceImport() {
  PsiFile bar = workspace.createPsiFile(WorkspacePath.createIfValid("foo/bar.py"));
  PsiFile source =
      workspace.createPsiFile(
          WorkspacePath.createIfValid("lib/source.py"), "from foo import bar");
  List<PyFromImportStatement> imports = ((PyFile) source).getFromImports();
  assertThat(imports).hasSize(1);
  assertThat(imports.get(0).getImportElements()[0].resolve()).isEqualTo(bar);
}
 

@Test
public void testResolveInitPy() {
  PsiFile initPy =
      workspace.createPsiFile(WorkspacePath.createIfValid("pyglib/flags/__init__.py"));
  PsiFile source =
      workspace.createPsiFile(
          WorkspacePath.createIfValid("lib/source.py"), "from pyglib import flags");
  List<PyFromImportStatement> imports = ((PyFile) source).getFromImports();
  assertThat(imports).hasSize(1);
  assertThat(imports.get(0).getImportElements()[0].resolve()).isEqualTo(initPy);
}
 

@Test
public void testResolveGenfiles() {
  BlazeInfo roots =
      BlazeProjectDataManager.getInstance(getProject()).getBlazeProjectData().getBlazeInfo();
  PsiFile genfile =
      fileSystem.createPsiFile(new File(roots.getGenfilesDirectory(), "foo/bar.py").getPath());
  PsiFile source =
      workspace.createPsiFile(
          WorkspacePath.createIfValid("lib/source.py"), "from foo import bar");
  List<PyFromImportStatement> imports = ((PyFile) source).getFromImports();
  assertThat(imports).hasSize(1);
  assertThat(imports.get(0).getImportElements()[0].resolve()).isEqualTo(genfile);
}
 

private boolean isTopLevelExpression(@NotNull CompletionParameters parameters) {
  PsiElement position = parameters.getPosition();
  PsiElement expression = position.getParent();
  if (!(expression instanceof PyExpression)) return false;
  PsiElement statement = expression.getParent();
  if (!(statement instanceof PyStatement)) return false;
  PsiElement file = statement.getParent();
  return file instanceof PyFile;
}
 

@NotNull
public static Map<String, PyReferenceExpression> findTargetDefinitions(@NotNull PyFile pyFile) {
  final PyFunction buildFileAliases = pyFile.findTopLevelFunction("build_file_aliases");
  final PyStatement[] statements =
    buildFileAliases != null ? buildFileAliases.getStatementList().getStatements() : PyStatement.EMPTY_ARRAY;
  final Map<String, PyReferenceExpression> result = new HashMap<>();
  for (PyStatement statement : statements) {
    if (!(statement instanceof PyReturnStatement)) {
      continue;
    }
    final PyExpression returnExpression = ((PyReturnStatement)statement).getExpression();
    if (!(returnExpression instanceof PyCallExpression)) {
      continue;
    }
    final PyArgumentList argumentList = ((PyCallExpression)returnExpression).getArgumentList();
    final Collection<PyKeywordArgument> targetDefinitions = PsiTreeUtil.findChildrenOfType(argumentList, PyKeywordArgument.class);
    for (PyKeywordArgument targets : targetDefinitions) {
      final PyExpression targetsExpression = targets != null ? targets.getValueExpression() : null;
      if (targetsExpression instanceof PyDictLiteralExpression) {
        for (PyKeyValueExpression keyValueExpression : ((PyDictLiteralExpression)targetsExpression).getElements()) {
          final PyExpression keyExpression = keyValueExpression.getKey();
          final PyExpression valueExpression = keyValueExpression.getValue();
          if (keyExpression instanceof PyStringLiteralExpression) {
            result.put(
              ((PyStringLiteralExpression)keyExpression).getStringValue(),
              valueExpression instanceof PyReferenceExpression ? (PyReferenceExpression)valueExpression : null
            );
          }
        }
      }
    }
  }
  return result;
}
 

public void testAliases() {
  final PsiFile psiFile = myFixture.configureByFile("aliases.py");
  assertTrue(psiFile instanceof PyFile);
  final Map<String, PyReferenceExpression> definitions = PantsPsiUtil.findTargetDefinitions((PyFile)psiFile);
  assertContainsElements(
    definitions.keySet(),
    Arrays.asList("annotation_processor", "jar_library", "scala_library", "scalac_plugin", "rglobs")
  );
}
 

private static boolean isPyFileOutsideProject(PsiElement element) {
  PsiFile file = element.getContainingFile();
  return file instanceof PyFile
      && !inProjectScope(file.getProject(), file.getViewProvider().getVirtualFile());
}
 

private TestLocation(
    PyFile testFile, @Nullable PyClass testClass, @Nullable PyFunction testFunction) {
  this.testFile = testFile;
  this.testClass = testClass;
  this.testFunction = testFunction;
}