path: root/python/src/com/jetbrains/python/inspections/quickfix/GenerateBinaryStubsFix.java

/*
 * Copyright 2000-2013 JetBrains s.r.o.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.jetbrains.python.inspections.quickfix;

import com.google.common.collect.Lists;
import com.intellij.codeInspection.LocalQuickFix;
import com.intellij.codeInspection.ProblemDescriptor;
import com.intellij.execution.process.ProcessOutput;
import com.intellij.openapi.application.ReadAction;
import com.intellij.openapi.application.Result;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.module.Module;
import com.intellij.openapi.module.ModuleUtilCore;
import com.intellij.openapi.progress.ProgressIndicator;
import com.intellij.openapi.progress.ProgressManager;
import com.intellij.openapi.progress.Task;
import com.intellij.openapi.progress.Task.Backgroundable;
import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.projectRoots.Sdk;
import com.intellij.openapi.vfs.LocalFileSystem;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.util.QualifiedName;
import com.intellij.util.Consumer;
import com.jetbrains.python.PyBundle;
import com.jetbrains.python.PythonHelpersLocator;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.sdk.InvalidSdkException;
import com.jetbrains.python.sdk.PySdkUtil;
import com.jetbrains.python.sdk.PythonSdkType;
import com.jetbrains.python.sdk.flavors.IronPythonSdkFlavor;
import com.jetbrains.python.sdk.flavors.PythonSdkFlavor;
import com.jetbrains.python.sdk.skeletons.PySkeletonGenerator;
import com.jetbrains.python.sdk.skeletons.PySkeletonRefresher;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.File;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

/**
 * @author yole
 */
public class GenerateBinaryStubsFix implements LocalQuickFix {
  private static final Logger LOG = Logger.getInstance("#" + GenerateBinaryStubsFix.class.getName());

  private final String myQualifiedName;
  private final Sdk mySdk;

  /**
   * Generates pack of fixes available for some unresolved import statement.
   * Be sure to call {@link #isApplicable(com.jetbrains.python.psi.PyImportStatementBase)} first to make sure this statement is supported
   *
   * @param importStatementBase statement to fix
   * @return pack of fixes
   */
  @NotNull
  public static Collection<GenerateBinaryStubsFix> generateFixes(@NotNull final PyImportStatementBase importStatementBase) {
    final List<String> names = importStatementBase.getFullyQualifiedObjectNames();
    final List<GenerateBinaryStubsFix> result = new ArrayList<GenerateBinaryStubsFix>(names.size());
    for (final String qualifiedName : names) {
      result.add(new GenerateBinaryStubsFix(importStatementBase, qualifiedName));
    }
    return result;
  }

  /**
   * @param importStatementBase statement to fix
   * @param qualifiedName       name should be fixed (one of {@link com.jetbrains.python.psi.PyImportStatementBase#getFullyQualifiedObjectNames()})
   */
  private GenerateBinaryStubsFix(@NotNull final PyImportStatementBase importStatementBase, @NotNull final String qualifiedName) {
    myQualifiedName = qualifiedName;
    mySdk = getPythonSdk(importStatementBase);
  }

  @Override
  @NotNull
  public String getName() {
    return PyBundle.message("sdk.gen.stubs.for.binary.modules", myQualifiedName);
  }

  @Override
  @NotNull
  public String getFamilyName() {
    return "Generate binary stubs";
  }

  @Override
  public void applyFix(@NotNull final Project project, @NotNull final ProblemDescriptor descriptor) {
    final PsiFile file = descriptor.getPsiElement().getContainingFile();
    final Backgroundable backgroundable = getFixTask(file);
    ProgressManager.getInstance().runProcessWithProgressAsynchronously(backgroundable, new BackgroundableProcessIndicator(backgroundable));
  }


  /**
   * Returns fix task that is used to generate stubs
   * @param fileToRunTaskIn file where task should run
   * @return task itself
   */
  @NotNull
  public Backgroundable getFixTask(@NotNull final PsiFile fileToRunTaskIn) {
    final Project project = fileToRunTaskIn.getProject();
    final String folder = fileToRunTaskIn.getContainingDirectory().getVirtualFile().getCanonicalPath();
    return new Task.Backgroundable(project, "Generating skeletons for binary module", false) {

      @Override
      public void run(@NotNull ProgressIndicator indicator) {
        indicator.setIndeterminate(true);


        final List<String> assemblyRefs = new ReadAction<List<String>>() {
          @Override
          protected void run(@NotNull Result<List<String>> result) throws Throwable {
            result.setResult(collectAssemblyReferences(fileToRunTaskIn));
          }
        }.execute().getResultObject();


        try {
          final PySkeletonRefresher refresher = new PySkeletonRefresher(project, null, mySdk, null, null, folder);

          if (needBinaryList(myQualifiedName)) {
            if (!generateSkeletonsForList(refresher, indicator, folder)) return;
          }
          else {
            //noinspection unchecked
            refresher.generateSkeleton(myQualifiedName, "", assemblyRefs, Consumer.EMPTY_CONSUMER);
          }
          final VirtualFile skeletonDir;
          skeletonDir = LocalFileSystem.getInstance().findFileByPath(refresher.getSkeletonsPath());
          if (skeletonDir != null) {
            skeletonDir.refresh(true, true);
          }
        }
        catch (InvalidSdkException e) {
          LOG.error(e);
        }
      }
    };
  }

  private boolean generateSkeletonsForList(@NotNull final PySkeletonRefresher refresher,
                                           ProgressIndicator indicator,
                                           @Nullable final String currentBinaryFilesPath) throws InvalidSdkException {
    final PySkeletonGenerator generator = new PySkeletonGenerator(refresher.getSkeletonsPath(), mySdk, currentBinaryFilesPath);
    indicator.setIndeterminate(false);
    final String homePath = mySdk.getHomePath();
    if (homePath == null) return false;
    final ProcessOutput runResult = PySdkUtil.getProcessOutput(
      new File(homePath).getParent(),
      new String[]{
        homePath,
        PythonHelpersLocator.getHelperPath("extra_syspath.py"), myQualifiedName},
      PythonSdkType.getVirtualEnvExtraEnv(homePath), 5000
    );
    if (runResult.getExitCode() == 0 && !runResult.isTimeout()) {
      final String extraPath = runResult.getStdout();
      final PySkeletonGenerator.ListBinariesResult binaries = generator.listBinaries(mySdk, extraPath);
      final List<String> names = Lists.newArrayList(binaries.modules.keySet());
      Collections.sort(names);
      final int size = names.size();
      for (int i = 0; i != size; ++i) {
        final String name = names.get(i);
        indicator.setFraction((double)i / size);
        if (needBinaryList(name)) {
          indicator.setText2(name);
          //noinspection unchecked
          refresher.generateSkeleton(name, "", new ArrayList<String>(), Consumer.EMPTY_CONSUMER);
        }
      }
    }
    return true;
  }

  private static boolean needBinaryList(@NotNull final String qualifiedName) {
    return qualifiedName.startsWith("gi.repository");
  }

  private List<String> collectAssemblyReferences(PsiFile file) {
    if (!(PythonSdkFlavor.getFlavor(mySdk) instanceof IronPythonSdkFlavor)) {
      return Collections.emptyList();
    }
    final List<String> result = new ArrayList<String>();
    file.accept(new PyRecursiveElementVisitor() {
      @Override
      public void visitPyCallExpression(PyCallExpression node) {
        super.visitPyCallExpression(node);
        // TODO: What if user loads it not by literal? We need to ask user for list of DLLs
        if (node.isCalleeText("AddReference", "AddReferenceByPartialName", "AddReferenceByName")) {
          final PyExpression[] args = node.getArguments();
          if (args.length == 1 && args[0] instanceof PyStringLiteralExpression) {
            result.add(((PyStringLiteralExpression)args[0]).getStringValue());
          }
        }
      }
    });
    return result;
  }

  /**
   * Checks if this fix can help you to generate binary stubs
   *
   * @param importStatementBase statement to fix
   * @return true if this fix could work
   */
  public static boolean isApplicable(@NotNull final PyImportStatementBase importStatementBase) {
    if (importStatementBase.getFullyQualifiedObjectNames().isEmpty()) {
      return false;
    }
    final Sdk sdk = getPythonSdk(importStatementBase);
    if (sdk == null) {
      return false;
    }
    final PythonSdkFlavor flavor = PythonSdkFlavor.getFlavor(sdk);
    if (flavor instanceof IronPythonSdkFlavor) {
      return true;
    }
    return isGtk(importStatementBase);
  }

  private static boolean isGtk(@NotNull final PyImportStatementBase importStatementBase) {
    if (importStatementBase instanceof PyFromImportStatement) {
      final QualifiedName qName = ((PyFromImportStatement)importStatementBase).getImportSourceQName();
      if (qName != null && qName.matches("gi", "repository")) {
        return true;
      }
    }
    return false;
  }

  @Nullable
  private static Sdk getPythonSdk(@NotNull final PsiElement element) {
    final Module module = ModuleUtilCore.findModuleForPsiElement(element);
    return (module == null) ? null : PythonSdkType.findPythonSdk(module);
  }
}