zig-cli/cli.zig

const std = @import("std");
const io = std.io;
const mem = std.mem;
const process = std.process;
const ArrayList = std.ArrayList;
const StringHashMap = std.StringHashMap;
const out = io.getStdOut().writer();
const util = @import("./util.zig");

pub const Option = struct {
    short: []const u8,
    long: []const u8,
    optional: bool = true,
    global: bool = false,
    description: []const u8,
    value_string: []const u8 = "",
    value_int: i32 = 0,
    value_float: f32 = 0.0,
    value_bool: bool = false,
    value_type: ValueType = .BOOL,

    const ValueType = enum {
        STRING,
        BOOL,
        INT,
        FLOAT,
    };
};

const help_option = Option{
    .short = "h",
    .long = "help",
    .global = true,
    .description = "Print help",
};

const version_option = Option{
    .short = "v",
    .long = "version",
    .global = true,
    .description = "Print version",
};

pub const Command = struct {
    const Commands = StringHashMap(Command);
    const Options = StringHashMap(Option);
    const Args = ArrayList([]const u8);
    pub const CommandConfig = struct {
        name: []const u8,
        short_description: []const u8,
        long_description: []const u8,
        usage: []const u8,
        version: []const u8 = "0.0.1",
        run: ?*const fn (cmd: *Command) anyerror!void,
    };

    subcommands: Commands,
    options: Options,
    arguments: Args,
    cmd_args: Args,
    allocator: mem.Allocator,
    config: CommandConfig,

    pub fn init(allocator: mem.Allocator, args: CommandConfig) !Command {
        var options = Options.init(allocator);
        try options.put(help_option.long, help_option);
        try options.put(version_option.long, version_option);
        return Command{
            .subcommands = Commands.init(allocator),
            .options = options,
            .arguments = Args.init(allocator),
            .cmd_args = Args.init(allocator),
            .allocator = allocator,
            .config = args,
        };
    }

    pub fn deinit(self: *Command) void {
        self.subcommands.deinit();
        self.options.deinit();
        self.arguments.deinit();
        self.cmd_args.deinit();
    }

    pub fn addOption(self: *Command, option: Option) !void {
        try self.options.put(option.long, option);
    }

    pub fn addCommand(self: *Command, command: Command) !void {
        var option_iterator = self.options.iterator();
        while (option_iterator.next()) |entry| {
            const option = entry.value_ptr.*;
            if (option.global) {
                try @constCast(&command).addOption(option);
            }
        }
        try self.subcommands.put(command.config.name, command);
    }

    pub fn help(self: *Command) !void {
        const description = if (self.config.long_description.len > 0)
            self.config.long_description
        else if (self.config.short_description.len > 0)
            self.config.short_description
        else
            self.config.name;

        try out.print("{s}\n\nUsage: {s}", .{ description, self.config.usage });
        if (self.subcommands.count() > 0) {
            try out.print(" [COMMAND]", .{});
        }
        try out.print(" [OPTIONS] [ARGUMENTS]\n\n", .{});

        if (self.subcommands.count() > 0) {
            try out.print("Commands: \n\n", .{});
            var entry_iterator = self.subcommands.iterator();
            while (entry_iterator.next()) |subcommand| {
                try out.print("    {s: <20}{s}\n", .{ subcommand.key_ptr.*, subcommand.value_ptr.*.config.short_description });
            }
            try out.print("\n", .{});
        }

        try out.print("Options: \n\n", .{});

        if (self.options.count() > 0) {
            var entry_iterator = self.options.iterator();
            while (entry_iterator.next()) |entry| {
                const option = entry.value_ptr.*;
                const flags = try mem.concat(self.allocator, u8, &[_][]const u8{ "-", option.short, ", --", option.long });
                try out.print("    {s: <20}{s}", .{ flags, option.description });
                if (!option.optional) {
                    try out.print(" (required)\n", .{});
                } else {
                    try out.print("\n", .{});
                }
            }
        }
    }

    pub fn printVersion(self: *Command) !void {
        _ = try io.getStdOut().write(try mem.concat(self.allocator, u8, &[_][]const u8{ self.config.version, "\n" }));
    }

    fn defaultRun(self: *Command) !void {
        if (self.options.get("help").?.value_bool) {
            try self.help();
        } else if (self.options.get("version").?.value_bool) {
            try self.printVersion();
        } else if (self.config.run) |run| {
            try run(self);
        }
    }

    fn parse(self: *Command) !void {
        var i: usize = 1;
        var current_command = self.*;
        var current_option: Option = undefined;
        var parsing_option = false;
        var subcommand_parse_over = false;
        var option_parse_over = false;
        parse_while_blk: while (i < self.cmd_args.items.len) {
            const current_arg = self.cmd_args.items[i];
            if (!subcommand_parse_over) {
                if (current_command.subcommands.get(current_arg)) |subcommand| {
                    current_command = subcommand;
                    i += 1;
                    continue :parse_while_blk;
                }
                subcommand_parse_over = true;
            }
            if (!option_parse_over) {
                if (parsing_option) {
                    switch (current_option.value_type) {
                        .STRING => current_option.value_string = current_arg,
                        .INT => current_option.value_int = util.strToInt(current_arg.ptr),
                        .FLOAT => current_option.value_float = util.strToFloat(current_arg.ptr),
                        else => {},
                    }
                    try current_command.options.put(current_option.long, current_option);
                    parsing_option = false;
                    i += 1;
                    continue :parse_while_blk;
                }
                var entry_iterator = current_command.options.iterator();
                while (entry_iterator.next()) |entry| {
                    const option = entry.value_ptr.*;
                    if (mem.eql(u8, current_arg, try mem.concat(current_command.allocator, u8, &[_][]const u8{ "-", option.short })) or
                        mem.eql(u8, current_arg, try mem.concat(current_command.allocator, u8, &[_][]const u8{ "--", option.long })))
                    {
                        current_option = option;
                        parsing_option = true;
                        if (current_option.value_type == .BOOL) {
                            current_option.value_bool = true;
                            try current_command.options.put(entry.key_ptr.*, current_option);
                            parsing_option = false;
                        }
                        i += 1;
                        continue :parse_while_blk;
                    }
                }
                option_parse_over = true;
            }
            try current_command.arguments.append(current_arg);
            i += 1;
        }
        try current_command.defaultRun();
    }

    pub fn execute(self: *Command) !void {
        var arg_iterator = process.args();
        while (arg_iterator.next()) |arg| {
            try self.cmd_args.append(arg);
        }
        self.cmd_args.items[0] = self.config.name;
        try self.parse();
    }

    pub fn getBoolOption(self: *Command, option_name: []const u8) bool {
        return if (self.options.get(option_name)) |option| option.value_bool else false;
    }

    pub fn getStringOption(self: *Command, option_name: []const u8) []const u8 {
        return if (self.options.get(option_name)) |option| option.value_string else "";
    }

    pub fn getIntOption(self: *Command, option_name: []const u8) i32 {
        return if (self.options.get(option_name)) |option| option.value_int else 0;
    }

    pub fn getFloatOption(self: *Command, option_name: []const u8) f32 {
        return if (self.options.get(option_name)) |option| option.value_float else 0.0;
    }
};