path: root/python/src/com/jetbrains/python/sdk/PySdkUtil.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.sdk;

import com.intellij.execution.ExecutionException;
import com.intellij.execution.process.CapturingProcessHandler;
import com.intellij.execution.process.ProcessOutput;
import com.intellij.execution.util.ExecUtil;
import com.intellij.openapi.application.PathManager;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.module.Module;
import com.intellij.openapi.projectRoots.Sdk;
import com.intellij.openapi.roots.OrderRootType;
import com.intellij.openapi.util.SystemInfo;
import com.intellij.openapi.util.io.FileUtil;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.openapi.vfs.VfsUtilCore;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.remote.RemoteSdkAdditionalData;
import com.intellij.util.SystemProperties;
import com.intellij.util.containers.HashMap;
import com.jetbrains.python.packaging.PyPackageUtil;
import com.jetbrains.python.packaging.PyRequirement;
import org.jetbrains.annotations.NonNls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * A more flexible cousin of SdkVersionUtil.
 * Needs not to be instantiated and only holds static methods.
 *
 * @author dcheryasov
 *         Date: Apr 24, 2008
 *         Time: 1:19:47 PM
 */
public class PySdkUtil {
  protected static final Logger LOG = Logger.getInstance("#com.jetbrains.python.sdk.SdkVersionUtil");

  // Windows EOF marker, Ctrl+Z
  public static final int SUBSTITUTE = 26;
  public static final String PATH_ENV_VARIABLE = "PATH";

  private PySdkUtil() {
    // explicitly none
  }

  /**
   * Executes a process and returns its stdout and stderr outputs as lists of lines.
   *
   * @param homePath process run directory
   * @param command  command to execute and its arguments
   * @return a tuple of (stdout lines, stderr lines, exit_code), lines in them have line terminators stripped, or may be null.
   */
  @NotNull
  public static ProcessOutput getProcessOutput(String homePath, @NonNls String[] command) {
    return getProcessOutput(homePath, command, -1);
  }

  /**
   * Executes a process and returns its stdout and stderr outputs as lists of lines.
   * Waits for process for possibly limited duration.
   *
   * @param homePath process run directory
   * @param command  command to execute and its arguments
   * @param timeout  how many milliseconds to wait until the process terminates; non-positive means inifinity.
   * @return a tuple of (stdout lines, stderr lines, exit_code), lines in them have line terminators stripped, or may be null.
   */
  @NotNull
  public static ProcessOutput getProcessOutput(String homePath, @NonNls String[] command, final int timeout) {
    return getProcessOutput(homePath, command, null, timeout);
  }

  @NotNull
  public static ProcessOutput getProcessOutput(String homePath,
                                               @NonNls String[] command,
                                               @Nullable @NonNls Map<String, String> extraEnv,
                                               final int timeout) {
    return getProcessOutput(homePath, command, extraEnv, timeout, null, true);
  }

  @NotNull
  public static ProcessOutput getProcessOutput(String homePath,
                                               @NonNls String[] command,
                                               @Nullable @NonNls Map<String, String> extraEnv,
                                               final int timeout,
                                               @Nullable byte[] stdin,
                                               boolean needEOFMarker) {
    if (homePath == null || !new File(homePath).exists()) {
      return new ProcessOutput();
    }
    final Map<String, String> systemEnv = System.getenv();
    final Map<String, String> env = extraEnv != null ? mergeEnvVariables(systemEnv, extraEnv) : systemEnv;
    try {
      final Process process = ExecUtil.exec(Arrays.asList(command), homePath, env);
      final CapturingProcessHandler processHandler = new CapturingProcessHandler(process);
      if (stdin != null) {
        final OutputStream processInput = processHandler.getProcessInput();
        assert processInput != null;
        processInput.write(stdin);
        if (SystemInfo.isWindows && needEOFMarker) {
          processInput.write(SUBSTITUTE);
          processInput.flush();
        }
        else {
          processInput.close();
        }
      }
      return processHandler.runProcess(timeout);
    }
    catch (ExecutionException e) {
      return getOutputForException(e);
    }
    catch (IOException e) {
      return getOutputForException(e);
    }
  }

  private static ProcessOutput getOutputForException(final Exception e) {
    LOG.warn(e);
    return new ProcessOutput() {
      @Override
      public String getStderr() {
        String err = super.getStderr();
        if (!StringUtil.isEmpty(err)) {
          err += "\n" + e.getMessage();
        }
        else {
          err = e.getMessage();
        }
        return err;
      }
    };
  }

  @NotNull
  public static Map<String, String> mergeEnvVariables(@NotNull Map<String, String> environment,
                                                       @NotNull Map<String, String> extraEnvironment) {
    final Map<String, String> result = new HashMap<String, String>(environment);
    for (Map.Entry<String, String> entry : extraEnvironment.entrySet()) {
      if (PATH_ENV_VARIABLE.equals(entry.getKey()) && result.containsKey(PATH_ENV_VARIABLE)) {
        result.put(PATH_ENV_VARIABLE, result.get(PATH_ENV_VARIABLE) + File.pathSeparator + entry.getValue());
      }
      else {
        result.put(entry.getKey(), entry.getValue());
      }
    }
    return result;
  }

  public static boolean isRemote(@Nullable Sdk sdk) {
    return sdk != null && sdk.getSdkAdditionalData() instanceof RemoteSdkAdditionalData;
  }

  public static String getUserSite() {
    if (SystemInfo.isWindows) {
      final String appdata = System.getenv("APPDATA");
      return appdata + File.separator + "Python";
    }
    else {
      final String userHome = SystemProperties.getUserHome();
      return userHome + File.separator + ".local";
    }
  }

  public static boolean isElementInSkeletons(@NotNull final PsiElement element) {
    final PsiFile file = element.getContainingFile();
    if (file != null) {
      final VirtualFile virtualFile = file.getVirtualFile();
      if (virtualFile != null) {
        final Sdk sdk = PythonSdkType.getSdk(element);
        if (sdk != null) {
          final VirtualFile skeletonsDir = findSkeletonsDir(sdk);
          if (skeletonsDir != null && VfsUtilCore.isAncestor(skeletonsDir, virtualFile, false)) {
            return true;
          }
        }
      }
    }
    return false;
  }

  public static String getRemoteSourcesLocalPath(String sdkHome) {
    String sep = File.separator;

    String basePath = PathManager.getSystemPath();
    return basePath +
           File.separator +
           PythonSdkType.REMOTE_SOURCES_DIR_NAME +
           sep +
           FileUtil.toSystemIndependentName(sdkHome).hashCode() +
           sep;
  }

  @Nullable
  public static VirtualFile findSkeletonsDir(@NotNull final Sdk sdk) {
    return findLibraryDir(sdk, PythonSdkType.SKELETON_DIR_NAME, PythonSdkType.BUILTIN_ROOT_TYPE);
  }

  @Nullable
  public static VirtualFile findAnyRemoteLibrary(@NotNull final Sdk sdk) {
    return findLibraryDir(sdk, PythonSdkType.REMOTE_SOURCES_DIR_NAME, OrderRootType.CLASSES);
  }

  private static VirtualFile findLibraryDir(Sdk sdk, String dirName, OrderRootType rootType) {
    final VirtualFile[] virtualFiles = sdk.getRootProvider().getFiles(rootType);
    for (VirtualFile virtualFile : virtualFiles) {
      if (virtualFile.isValid() && virtualFile.getPath().contains(dirName)) {
        return virtualFile;
      }
    }
    return null;
  }

  @Nullable
  public static List<PyRequirement> getRequirementsFromTxt(Module module) {
    final VirtualFile requirementsTxt = PyPackageUtil.findRequirementsTxt(module);
    if (requirementsTxt != null) {
      return PyRequirement.parse(requirementsTxt);
    }
    return null;
  }
}