org.apache.zeppelin.python.PythonCondaInterpreter.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.zeppelin.python.PythonCondaInterpreter.java

Source

/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements.  See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You 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 org.apache.zeppelin.python;

import org.apache.commons.lang.StringUtils;
import org.apache.zeppelin.interpreter.*;
import org.apache.zeppelin.interpreter.InterpreterResult.Code;
import org.apache.zeppelin.interpreter.InterpreterResult.Type;
import org.apache.zeppelin.scheduler.Scheduler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Conda support
 */
public class PythonCondaInterpreter extends Interpreter {
    Logger logger = LoggerFactory.getLogger(PythonCondaInterpreter.class);
    public static final String ZEPPELIN_PYTHON = "zeppelin.python";
    public static final String CONDA_PYTHON_PATH = "/bin/python";
    public static final String DEFAULT_ZEPPELIN_PYTHON = "python";

    public static final Pattern PATTERN_OUTPUT_ENV_LIST = Pattern.compile("([^\\s]*)[\\s*]*\\s(.*)");
    public static final Pattern PATTERN_COMMAND_ENV_LIST = Pattern.compile("env\\s*list\\s?");
    public static final Pattern PATTERN_COMMAND_ENV = Pattern.compile("env\\s*(.*)");
    public static final Pattern PATTERN_COMMAND_LIST = Pattern.compile("list");
    public static final Pattern PATTERN_COMMAND_CREATE = Pattern.compile("create\\s*(.*)");
    public static final Pattern PATTERN_COMMAND_ACTIVATE = Pattern.compile("activate\\s*(.*)");
    public static final Pattern PATTERN_COMMAND_DEACTIVATE = Pattern.compile("deactivate");
    public static final Pattern PATTERN_COMMAND_INSTALL = Pattern.compile("install\\s*(.*)");
    public static final Pattern PATTERN_COMMAND_UNINSTALL = Pattern.compile("uninstall\\s*(.*)");
    public static final Pattern PATTERN_COMMAND_HELP = Pattern.compile("help");
    public static final Pattern PATTERN_COMMAND_INFO = Pattern.compile("info");

    public PythonCondaInterpreter(Properties property) {
        super(property);
    }

    @Override
    public void open() {

    }

    @Override
    public void close() {

    }

    @Override
    public InterpreterResult interpret(String st, InterpreterContext context) {
        InterpreterOutput out = context.out;
        Matcher activateMatcher = PATTERN_COMMAND_ACTIVATE.matcher(st);
        Matcher createMatcher = PATTERN_COMMAND_CREATE.matcher(st);
        Matcher installMatcher = PATTERN_COMMAND_INSTALL.matcher(st);
        Matcher uninstallMatcher = PATTERN_COMMAND_UNINSTALL.matcher(st);
        Matcher envMatcher = PATTERN_COMMAND_ENV.matcher(st);

        try {
            if (PATTERN_COMMAND_ENV_LIST.matcher(st).matches()) {
                String result = runCondaEnvList();
                return new InterpreterResult(Code.SUCCESS, Type.HTML, result);
            } else if (envMatcher.matches()) {
                // `envMatcher` should be used after `listEnvMatcher`
                String result = runCondaEnv(getRestArgsFromMatcher(envMatcher));
                return new InterpreterResult(Code.SUCCESS, Type.HTML, result);
            } else if (PATTERN_COMMAND_LIST.matcher(st).matches()) {
                String result = runCondaList();
                return new InterpreterResult(Code.SUCCESS, Type.HTML, result);
            } else if (createMatcher.matches()) {
                String result = runCondaCreate(getRestArgsFromMatcher(createMatcher));
                return new InterpreterResult(Code.SUCCESS, Type.HTML, result);
            } else if (activateMatcher.matches()) {
                String envName = activateMatcher.group(1).trim();
                return runCondaActivate(envName);
            } else if (PATTERN_COMMAND_DEACTIVATE.matcher(st).matches()) {
                return runCondaDeactivate();
            } else if (installMatcher.matches()) {
                String result = runCondaInstall(getRestArgsFromMatcher(installMatcher));
                return new InterpreterResult(Code.SUCCESS, Type.HTML, result);
            } else if (uninstallMatcher.matches()) {
                String result = runCondaUninstall(getRestArgsFromMatcher(uninstallMatcher));
                return new InterpreterResult(Code.SUCCESS, Type.HTML, result);
            } else if (st == null || PATTERN_COMMAND_HELP.matcher(st).matches()) {
                runCondaHelp(out);
                return new InterpreterResult(Code.SUCCESS);
            } else if (PATTERN_COMMAND_INFO.matcher(st).matches()) {
                String result = runCondaInfo();
                return new InterpreterResult(Code.SUCCESS, Type.HTML, result);
            } else {
                return new InterpreterResult(Code.ERROR, "Not supported command: " + st);
            }
        } catch (RuntimeException | IOException | InterruptedException e) {
            throw new InterpreterException(e);
        }
    }

    private void changePythonEnvironment(String envName) throws IOException, InterruptedException {
        PythonInterpreter python = getPythonInterpreter();
        String binPath = null;
        if (envName == null) {
            binPath = getProperty(ZEPPELIN_PYTHON);
            if (binPath == null) {
                binPath = DEFAULT_ZEPPELIN_PYTHON;
            }
        } else {
            Map<String, String> envList = getCondaEnvs();
            for (String name : envList.keySet()) {
                if (envName.equals(name)) {
                    binPath = envList.get(name) + CONDA_PYTHON_PATH;
                    break;
                }
            }
        }
        python.setPythonCommand(binPath);
    }

    private void restartPythonProcess() {
        PythonInterpreter python = getPythonInterpreter();
        python.close();
        python.open();
    }

    protected PythonInterpreter getPythonInterpreter() {
        LazyOpenInterpreter lazy = null;
        PythonInterpreter python = null;
        Interpreter p = getInterpreterInTheSameSessionByClassName(PythonInterpreter.class.getName());

        while (p instanceof WrappedInterpreter) {
            if (p instanceof LazyOpenInterpreter) {
                lazy = (LazyOpenInterpreter) p;
            }
            p = ((WrappedInterpreter) p).getInnerInterpreter();
        }
        python = (PythonInterpreter) p;

        if (lazy != null) {
            lazy.open();
        }
        return python;
    }

    public static String runCondaCommandForTextOutput(String title, List<String> commands)
            throws IOException, InterruptedException {

        String result = runCommand(commands);
        return wrapCondaBasicOutputStyle(title, result);
    }

    private String runCondaCommandForTableOutput(String title, List<String> commands)
            throws IOException, InterruptedException {

        StringBuilder sb = new StringBuilder();
        String result = runCommand(commands);

        // use table output for pretty output
        Map<String, String> envPerName = parseCondaCommonStdout(result);
        return wrapCondaTableOutputStyle(title, envPerName);
    }

    protected Map<String, String> getCondaEnvs() throws IOException, InterruptedException {
        String result = runCommand("conda", "env", "list");
        Map<String, String> envList = parseCondaCommonStdout(result);
        return envList;
    }

    private String runCondaEnvList() throws IOException, InterruptedException {
        return wrapCondaTableOutputStyle("Environment List", getCondaEnvs());
    }

    private String runCondaEnv(List<String> restArgs) throws IOException, InterruptedException {

        restArgs.add(0, "conda");
        restArgs.add(1, "env");
        restArgs.add(3, "--yes"); // --yes should be inserted after command

        return runCondaCommandForTextOutput(null, restArgs);
    }

    private InterpreterResult runCondaActivate(String envName) throws IOException, InterruptedException {

        if (null == envName || envName.isEmpty()) {
            return new InterpreterResult(Code.ERROR, "Env name should be specified");
        }

        changePythonEnvironment(envName);
        restartPythonProcess();

        return new InterpreterResult(Code.SUCCESS, "'" + envName + "' is activated");
    }

    private InterpreterResult runCondaDeactivate() throws IOException, InterruptedException {

        changePythonEnvironment(null);
        restartPythonProcess();
        return new InterpreterResult(Code.SUCCESS, "Deactivated");
    }

    private String runCondaList() throws IOException, InterruptedException {
        List<String> commands = new ArrayList<String>();
        commands.add("conda");
        commands.add("list");

        return runCondaCommandForTableOutput("Installed Package List", commands);
    }

    private void runCondaHelp(InterpreterOutput out) {
        try {
            out.setType(InterpreterResult.Type.HTML);
            out.writeResource("output_templates/conda_usage.html");
        } catch (IOException e) {
            logger.error("Can't print usage", e);
        }
    }

    private String runCondaInfo() throws IOException, InterruptedException {
        List<String> commands = new ArrayList<String>();
        commands.add("conda");
        commands.add("info");

        return runCondaCommandForTextOutput("Conda Information", commands);
    }

    private String runCondaCreate(List<String> restArgs) throws IOException, InterruptedException {
        restArgs.add(0, "conda");
        restArgs.add(1, "create");
        restArgs.add(2, "--yes");

        return runCondaCommandForTextOutput("Environment Creation", restArgs);
    }

    private String runCondaInstall(List<String> restArgs) throws IOException, InterruptedException {

        restArgs.add(0, "conda");
        restArgs.add(1, "install");
        restArgs.add(2, "--yes");

        return runCondaCommandForTextOutput("Package Installation", restArgs);
    }

    private String runCondaUninstall(List<String> restArgs) throws IOException, InterruptedException {

        restArgs.add(0, "conda");
        restArgs.add(1, "uninstall");
        restArgs.add(2, "--yes");

        return runCondaCommandForTextOutput("Package Uninstallation", restArgs);
    }

    public static String wrapCondaBasicOutputStyle(String title, String content) {
        StringBuilder sb = new StringBuilder();
        if (null != title && !title.isEmpty()) {
            sb.append("<h4>").append(title).append("</h4>\n").append("</div><br />\n");
        }
        sb.append("<div style=\"white-space:pre-wrap;\">\n").append(content).append("</div>");

        return sb.toString();
    }

    public static String wrapCondaTableOutputStyle(String title, Map<String, String> kv) {
        StringBuilder sb = new StringBuilder();

        if (null != title && !title.isEmpty()) {
            sb.append("<h4>").append(title).append("</h4>\n");
        }

        sb.append("<div style=\"display:table;white-space:pre-wrap;\">\n");
        for (String name : kv.keySet()) {
            String path = kv.get(name);

            sb.append(String.format(
                    "<div style=\"display:table-row\">" + "<div style=\"display:table-cell;width:150px\">%s</div>"
                            + "<div style=\"display:table-cell;\">%s</div>" + "</div>\n",
                    name, path));
        }
        sb.append("</div>\n");

        return sb.toString();
    }

    public static Map<String, String> parseCondaCommonStdout(String out) throws IOException, InterruptedException {

        Map<String, String> kv = new LinkedHashMap<String, String>();
        String[] lines = out.split("\n");
        for (String s : lines) {
            if (s == null || s.isEmpty() || s.startsWith("#")) {
                continue;
            }
            Matcher match = PATTERN_OUTPUT_ENV_LIST.matcher(s);

            if (!match.matches()) {
                continue;
            }
            kv.put(match.group(1), match.group(2));
        }

        return kv;
    }

    @Override
    public void cancel(InterpreterContext context) {

    }

    @Override
    public FormType getFormType() {
        return FormType.NONE;
    }

    @Override
    public int getProgress(InterpreterContext context) {
        return 0;
    }

    /**
     * Use python interpreter's scheduler.
     * To make sure %python.conda paragraph and %python paragraph runs sequentially
     */
    @Override
    public Scheduler getScheduler() {
        PythonInterpreter pythonInterpreter = getPythonInterpreter();
        if (pythonInterpreter != null) {
            return pythonInterpreter.getScheduler();
        } else {
            return null;
        }
    }

    public static String runCommand(List<String> commands) throws IOException, InterruptedException {

        StringBuilder sb = new StringBuilder();

        ProcessBuilder builder = new ProcessBuilder(commands);
        builder.redirectErrorStream(true);
        Process process = builder.start();
        InputStream stdout = process.getInputStream();
        BufferedReader br = new BufferedReader(new InputStreamReader(stdout));
        String line;
        while ((line = br.readLine()) != null) {
            sb.append(line);
            sb.append("\n");
        }
        int r = process.waitFor(); // Let the process finish.

        if (r != 0) {
            throw new RuntimeException(
                    "Failed to execute `" + StringUtils.join(commands, " ") + "` exited with " + r);
        }

        return sb.toString();
    }

    public static String runCommand(String... command) throws IOException, InterruptedException {

        List<String> list = new ArrayList<>(command.length);
        for (String arg : command) {
            list.add(arg);
        }

        return runCommand(list);
    }

    public static List<String> getRestArgsFromMatcher(Matcher m) {
        // Arrays.asList just returns fixed-size, so we should use ctor instead of
        return new ArrayList<>(Arrays.asList(m.group(1).split(" ")));
    }
}