ArgParserImpl.java

/*
 * Copyright (c) 2016, Stein Eldar Johnsen
 *
 * 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 net.morimekta.terminal.args.impl;

import net.morimekta.collect.UnmodifiableMap;
import net.morimekta.strings.EscapeUtil;
import net.morimekta.terminal.args.Arg;
import net.morimekta.terminal.args.ArgException;
import net.morimekta.terminal.args.ArgNameFormat;
import net.morimekta.terminal.args.ArgParser;
import net.morimekta.terminal.args.Flag;
import net.morimekta.terminal.args.Option;
import net.morimekta.terminal.args.SubCommandSet;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Objects.requireNonNull;
import static net.morimekta.collect.UnmodifiableList.asList;
import static net.morimekta.collect.UnmodifiableList.listOf;

/**
 * Argument argumentParser class. This is the actual argumentParser that is initialized with
 * a set of options, arguments and properties, and is then initialized with
 * the appropriate fields.
 */
public class ArgParserImpl implements ArgParser {
    @Override
    public ArgParserImpl getParent() {
        return parent;
    }

    @Override
    public String getProgram() {
        return program;
    }

    @Override
    public String getVersion() {
        return version;
    }

    @Override
    public String getDescription() {
        return description;
    }

    @Override
    public List<Option> getOptions() {
        return options;
    }

    @Override
    public List<Arg> getArguments() {
        return arguments;
    }

    @SuppressWarnings("unchecked")
    public <SubCommandDef> SubCommandSet<SubCommandDef> getSubCommandSet() {
        return (SubCommandSet<SubCommandDef>) subCommandSet;
    }

    @Override
    public ArgParser parse(List<String> args) {
        try {
            parseInternal(asList(args));
        } catch (ArgException e) {
            if (e.getParser() == null) {
                throw new ArgException(this, e.getMessage(), e);
            }
            throw e;
        }
        return this;
    }

    @Override
    public void validate() {
        options.forEach(Arg::validate);
        arguments.forEach(Arg::validate);
        if (subCommandSet != null) {
            subCommandSet.validate();
        }
    }

    @Override
    public String toString() {
        return "ArgParser{" + getProgram() + "}";
    }

    public static class BuilderImpl
            extends ArgParserBuilderImpl
            implements ArgParser.Builder {
        public BuilderImpl(ArgParserImpl parent,
                           String program,
                           String version,
                           String description) {
            super(parent,
                  program,
                  version,
                  description,
                  new ArrayList<>(),
                  new ArrayList<>(),
                  new HashMap<>(),
                  new HashMap<>());
        }

        @Override
        public ArgParser build() {
            return new ArgParserImpl(
                    parent,
                    program,
                    version,
                    description,
                    options,
                    arguments,
                    longNameOptions,
                    shortOptions,
                    null);
        }

        /**
         * Add a command line option.
         *
         * @param arg The argument or option to add.
         * @param <A> The argument type.
         * @return The argument argumentParser.
         */
        @Override
        public <A extends Arg> ArgParser.Builder add(A arg) {
            if (arg instanceof Option) {
                Option option = (Option) arg;
                if (option.getName() != null) {
                    if (longNameOptions.containsKey(option.getName())) {
                        throw new IllegalArgumentException("Option " + option.getName() + " already exists");
                    }
                    if (parent != null && parent.getLongNameOption(option.getName()) != null) {
                        throw new IllegalArgumentException("Option " + option.getName() + " already exists in parent");
                    }

                    longNameOptions.put(option.getName(), option);
                }

                if (option instanceof Flag) {
                    String negate = ((Flag) option).getNegateName();
                    if (negate != null) {
                        if (longNameOptions.containsKey(negate)) {
                            throw new IllegalArgumentException("Flag " + negate + " already exists");
                        }
                        if (parent != null && parent.getLongNameOption(negate) != null) {
                            throw new IllegalArgumentException("Flag " + negate + " already exists in parent");
                        }

                        longNameOptions.put(negate, option);
                    }
                }

                if (option.getShortNames().length() > 0) {
                    for (char s : option.getShortNames().toCharArray()) {
                        if (shortOptions.containsKey(s)) {
                            throw new IllegalArgumentException("Short option -" + s + " already exists");
                        }
                        if (parent != null && parent.getShortNameOption(s) != null) {
                            throw new IllegalArgumentException("Short option -" + s + " already exists in parent");
                        }

                        shortOptions.put(s, option);
                    }
                }

                this.options.add(option);
            } else {
                // No arguments can be added after a sub-command-set.
                if (arguments.size() > 0 && arguments.get(arguments.size() - 1) instanceof SubCommandSetImpl) {
                    throw new IllegalArgumentException("No arguments can be added after a sub-command set");
                }

                arguments.add(arg);
            }
            return this;
        }

        /**
         * Add an argument builder.
         *
         * @param argBuilder The argument to add.
         * @param <A>        The base argument type.
         * @return The argument parser builder.
         */
        @Override
        public <A extends Arg> Builder add(Arg.Builder<A> argBuilder) {
            return this.add(requireNonNull(argBuilder, "argBuilder == null").build());
        }

        @Override
        public Builder generateArgsNameFormat(ArgNameFormat format) {
            this.defaultNameFormat = format;
            return this;
        }

        /**
         * Generate arguments to fill in the POJO config object.
         *
         * @param config The config object to fill in.
         * @return The argument parser builder.
         */
        @Override
        public Builder generateArgs(Object config) {
            return Generator.generateArgs(
                    this,
                    requireNonNull(config, "config == null"),
                    "",
                    defaultNameFormat);
        }

        @Override
        public <NewDef> SubCommandSet.Builder<NewDef> withSubCommands(String name,
                                                                      String usage,
                                                                      Consumer<NewDef> consumer) {
            return new SubCommandSetImpl.BuilderImpl<>(
                    name,
                    usage,
                    consumer,
                    this.parent,
                    this.program,
                    this.version,
                    this.description,
                    this.options,
                    this.arguments,
                    this.longNameOptions,
                    this.shortOptions);
        }

        private ArgNameFormat defaultNameFormat = ArgNameFormat.LISP;
    }

    private void parseInternal(List<String> args) {
        while (args.size() > 0) {
            String cur = args.get(0);

            if (cur.equals("--")) {
                // The remaining *must* be arguments.
                args = args.subList(1, args.size());
                break;
            } else if (cur.startsWith("--")) {
                // long opt.
                if (cur.contains("=")) {
                    cur = cur.substring(0, cur.indexOf("="));
                }
                Option opt = getLongNameOption(cur);
                if (opt == null) {
                    if (parent != null) {
                        throw new ArgException("No option for %s on %s", cur, getProgram());
                    }
                    throw new ArgException("No option for %s", cur);
                }
                args = args.subList(opt.apply(args), args.size());
            } else if (cur.startsWith("-")) {
                // short opt.
                String remaining = cur.substring(1);
                int skip = 0;
                while (remaining.length() > 0) {
                    Option opt = getShortNameOption(remaining.charAt(0));
                    if (opt == null) {
                        if (parent != null) {
                            throw new ArgException("No short opt for %s on %s", remaining.charAt(0), getProgram());
                        }
                        throw new ArgException("No short opt for -%s", remaining.charAt(0));
                    }
                    skip = opt.applyShort(remaining, args);
                    if (skip == 0) {
                        remaining = remaining.substring(1);
                    } else {
                        break;
                    }
                }
                args = args.subList(Math.max(1, skip), args.size());
            } else {
                if (cur.startsWith("@")) {
                    File f = new File(cur.substring(1));
                    if (f.exists() && f.isFile()) {
                        try (FileInputStream fis = new FileInputStream(f);
                             BufferedReader reader = new BufferedReader(new InputStreamReader(fis, UTF_8))) {
                            List<String> lines = reader.lines()
                                                       .map(String::trim)
                                                       // strip empty lines and commented lines
                                                       .filter(l -> !(l.isEmpty() || l.startsWith("#")))
                                                       // be smart about splitting lines into single args.
                                                       .flatMap(this::argFileLineStream)
                                                       .collect(Collectors.toList());
                            if (lines.size() > 0) {
                                parse(asList(lines));
                            }
                        } catch (ArgException e) {
                            throw new ArgException(this, "Argument file %s: %s", f.getName(), e.getMessage(), e);
                        } catch (IOException e) {
                            throw new ArgException(this, e.getMessage(), e);
                        }
                        args = args.subList(1, args.size());
                        continue;
                    }
                }

                // Argument / sub-command.
                int consumed = 0;
                for (Arg arg : arguments) {
                    consumed = arg.apply(args);
                    if (consumed > 0) {
                        break;
                    }
                }
                if (consumed == 0 && subCommandSet != null) {
                    consumed = subCommandSet.apply(args);
                }
                if (consumed == 0) {
                    throw new ArgException("No option found for %s", args.get(0));
                }
                args = args.subList(consumed, args.size());
            }
        }

        if (args.isEmpty() &&
            subCommandSet != null &&
            !subCommandSet.isSelected() &&
            subCommandSet.getDefaultValue() != null) {
            // If no sub-command is selected by this point, and there is a
            // default subcommand, then use that with no extra options.
            subCommandSet.apply(listOf(subCommandSet.getDefaultValue()));
        }

        // Then consume the rest as arguments. Note that sub-commands may
        // *not* come after '--'.
        while (args.size() > 0) {
            // Argument / sub-command.
            int consumed = 0;
            for (Arg arg : arguments) {
                consumed = arg.apply(args);
                if (consumed > 0) {
                    break;
                }
            }
            if (consumed == 0) {
                throw new ArgException("No argument found for %s", args.get(0));
            }
            args = args.subList(consumed, args.size());
        }
    }

    private Stream<String> argFileLineStream(String line) {
        // If the line is double-quoted (entirely), treat the whole
        // as a json-escaped string argument.
        if (line.startsWith("\"") && line.endsWith("\"")) {
            return Stream.of(EscapeUtil.javaUnEscape(line.substring(1, line.length() - 1)));
        }
        // otherwise split on the first space,
        // and treat each part an arg entry, so that
        // "--something and more"
        // becomes:
        // ["--something", "and more"]
        return Arrays.stream(line.split(" ", 2))
                     .map(String::trim);
    }

    private Option getLongNameOption(String name) {
        Option option = longNameOptions.get(name);
        if (option == null && parent != null) {
            option = parent.getLongNameOption(name);
        }
        return option;
    }

    private Option getShortNameOption(char c) {
        Option option = shortOptions.get(c);
        if (option == null && parent != null) {
            option = parent.getShortNameOption(c);
        }
        return option;
    }

    private final ArgParserImpl parent;

    private final String program;
    private final String version;
    private final String description;

    private final List<Option>         options;
    private final List<Arg>            arguments;
    private final SubCommandSetImpl<?> subCommandSet;

    private final Map<String, Option>    longNameOptions;
    private final Map<Character, Option> shortOptions;

    protected ArgParserImpl(ArgParserImpl parent,
                            String program,
                            String version,
                            String description,
                            List<Option> options,
                            List<Arg> arguments,
                            Map<String, Option> longNameOptions,
                            Map<Character, Option> shortOptions,
                            SubCommandSetImpl<?> subCommandSet) {
        this.parent = parent;
        this.program = program;
        this.version = version;
        this.description = description;
        this.options = asList(options);
        this.arguments = asList(arguments);
        this.longNameOptions = UnmodifiableMap.asMap(longNameOptions);
        this.shortOptions = UnmodifiableMap.asMap(shortOptions);
        this.subCommandSet = subCommandSet;
    }
}