wyag

Write yourself a git
Log | Files | Refs | README

argparse.zig (20643B)


      1 const std = @import("std");
      2 
      3 const ArgParse = @This();
      4 
      5 pub const Subcommand = struct {
      6     parent: *ArgParse,
      7     name: []const u8,
      8     description: []const u8,
      9     positionals: std.ArrayListUnmanaged(*Positional) = .{},
     10     flags: std.ArrayListUnmanaged(*Flag) = .{},
     11 
     12     wasExecuted: bool = false,
     13 
     14     pub fn addPositional(self: *Subcommand, pos: *Positional) !void {
     15         try self.positionals.append(self.parent.aa.allocator(), pos);
     16     }
     17 
     18     pub fn addFlag(self: *Subcommand, flag: *Flag) !void {
     19         try self.flags.append(self.parent.aa.allocator(), flag);
     20     }
     21 };
     22 
     23 pub const Positional = struct {
     24     name: []const u8,
     25     description: []const u8,
     26     default: ?[]const u8 = null,
     27 
     28     value: ?[]const u8 = null,
     29 };
     30 
     31 pub const Flag = struct {
     32     long: []const u8,
     33     short: ?[]const u8 = null,
     34     description: []const u8,
     35     hasarg: bool,
     36 
     37     waspresent: bool = false,
     38     defaultargvalue: ?[]const u8 = null,
     39     argvalue: ?[]const u8 = null,
     40 };
     41 
     42 aa: std.heap.ArenaAllocator,
     43 progname: []const u8,
     44 description: []const u8,
     45 subcommands: std.ArrayListUnmanaged(*Subcommand) = .{},
     46 positionals: std.ArrayListUnmanaged(*Positional) = .{},
     47 flags: std.ArrayListUnmanaged(*Flag) = .{},
     48 excess: std.ArrayListUnmanaged([]const u8) = .{},
     49 
     50 pub fn init(ca: std.mem.Allocator, progname: []const u8, description: []const u8) ArgParse {
     51     return .{
     52         .aa = std.heap.ArenaAllocator.init(ca),
     53         .progname = progname,
     54         .description = description,
     55     };
     56 }
     57 
     58 pub fn deinit(self: *ArgParse) void {
     59     self.aa.deinit();
     60 }
     61 
     62 pub fn addSubcommand(self: *ArgParse, subcommand: *Subcommand) !void {
     63     try self.subcommands.append(self.aa.allocator(), subcommand);
     64 }
     65 
     66 pub fn addPositional(self: *ArgParse, pos: *Positional) !void {
     67     try self.positionals.append(self.aa.allocator(), pos);
     68 }
     69 
     70 pub fn addFlag(self: *ArgParse, flag: *Flag) !void {
     71     try self.flags.append(self.aa.allocator(), flag);
     72 }
     73 
     74 // Return false if args were not parsed (but it wasn't an error)
     75 pub fn parseOrHelp(self: *ArgParse) !bool {
     76     self.parse() catch |e| {
     77         switch (e) {
     78             error.HelpWanted => {
     79                 try self.help(std.io.getStdOut().writer());
     80                 return false;
     81             },
     82             error.InvalidArgs => {
     83                 try self.help(std.io.getStdErr().writer());
     84                 return e;
     85             },
     86             else => return e,
     87         }
     88     };
     89     return true;
     90 }
     91 
     92 pub fn parse(self: *ArgParse) !void {
     93     var it = try std.process.argsWithAllocator(self.aa.allocator());
     94     defer it.deinit();
     95     try self.parseInternal(&it);
     96 }
     97 
     98 pub fn help(self: *ArgParse, writer: anytype) !void {
     99     // Add the default help argument
    100     var helpflag: Flag = .{
    101         .long = "help",
    102         .short = "h",
    103         .description = "Show this help",
    104         .hasarg = false,
    105     };
    106     try self.flags.append(self.aa.allocator(), &helpflag);
    107 
    108     try std.fmt.format(writer, "Usage: {s}", .{self.progname});
    109     for (self.positionals.items) |pos| {
    110         try std.fmt.format(writer, " [{s}]", .{pos.name});
    111     }
    112     try std.fmt.format(writer, "\n{s}\n", .{self.description});
    113     for (self.positionals.items) |pos| {
    114         try std.fmt.format(writer, "  {s}: {s}", .{ pos.name, pos.description });
    115         if (pos.default) |df| {
    116             try std.fmt.format(writer, " (Default: {s})", .{df});
    117         }
    118         try writer.writeByte('\n');
    119     }
    120     try std.fmt.format(writer, "\nFlags:\n", .{});
    121     var longestFlag: u64 = 0;
    122     for (self.flags.items) |flag| {
    123         var cw = std.io.countingWriter(std.io.null_writer);
    124         try writeFlagsLongShort(flag, cw.writer());
    125         longestFlag = @max(longestFlag, cw.bytes_written);
    126     }
    127     longestFlag += 3;
    128     for (self.flags.items) |flag| {
    129         var cw = std.io.countingWriter(writer);
    130         const w = cw.writer();
    131         try w.writeByteNTimes(' ', 2);
    132         try writeFlagsLongShort(flag, w);
    133         const pad = (longestFlag - cw.bytes_written);
    134         try writer.writeByteNTimes(' ', pad);
    135         try std.fmt.format(writer, "{s}", .{flag.description});
    136         if (flag.defaultargvalue) |dav| {
    137             try std.fmt.format(writer, " (Default: {s})", .{dav});
    138         }
    139         try writer.writeByte('\n');
    140     }
    141 }
    142 
    143 fn writeFlagsLongShort(flag: *Flag, writer: anytype) !void {
    144     if (flag.short) |short| {
    145         try std.fmt.format(writer, "-{s},", .{short});
    146     }
    147     try std.fmt.format(writer, "--{s}", .{flag.long});
    148     if (flag.hasarg) {
    149         try std.fmt.format(writer, "=ARG", .{});
    150     }
    151 }
    152 
    153 fn parseInternal(self: *ArgParse, it: anytype) !void {
    154     var pos_i: usize = 0;
    155     var subcommand: ?*Subcommand = null;
    156     if (!it.skip()) return error.InvalidArgs;
    157     lp: while (it.next()) |nxt| {
    158         if (std.mem.startsWith(u8, nxt, "-")) {
    159             var flag = nxt[1..];
    160             const longflag = std.mem.startsWith(u8, nxt, "--");
    161             if (longflag) {
    162                 flag = nxt[2..];
    163             }
    164             var spl = std.mem.splitScalar(u8, flag, '=');
    165             const key = spl.first();
    166             if (std.mem.eql(u8, key, "help") or std.mem.eql(u8, key, "h")) {
    167                 return error.HelpWanted;
    168             }
    169 
    170             const flags: []*Flag = if (subcommand) |sc| sc.flags.items else self.flags.items;
    171             for (flags) |aflag| {
    172                 const ismatch = if (longflag) std.mem.eql(u8, key, aflag.long) else if (aflag.short) |srt| std.mem.eql(u8, key, srt) else false;
    173                 if (ismatch) {
    174                     if (aflag.waspresent) {
    175                         return error.InvalidArgs;
    176                     }
    177                     aflag.waspresent = true;
    178                     if (aflag.hasarg) {
    179                         if (spl.rest().len > 0) {
    180                             aflag.argvalue = try self.aa.allocator().dupe(u8, spl.rest());
    181                         } else {
    182                             const val = it.next() orelse return error.InvalidArgs;
    183                             aflag.argvalue = try self.aa.allocator().dupe(u8, val);
    184                         }
    185                     } else {
    186                         if (spl.rest().len > 0) {
    187                             return error.InvalidArgs;
    188                         }
    189                     }
    190                     break;
    191                 }
    192             } else {
    193                 return error.InvalidArgs;
    194             }
    195         } else {
    196             // Check for a matching subcommand to step into
    197             if (pos_i == 0 and subcommand == null) {
    198                 for (self.subcommands.items) |sc| {
    199                     if (std.mem.eql(u8, sc.name, nxt)) {
    200                         sc.wasExecuted = true;
    201                         subcommand = sc;
    202                         continue :lp;
    203                     }
    204                 }
    205             }
    206 
    207             const positionals: []*Positional = if (subcommand) |sc| sc.positionals.items else self.positionals.items;
    208             if (pos_i < positionals.len) {
    209                 positionals[pos_i].value = try self.aa.allocator().dupe(u8, nxt);
    210                 pos_i += 1;
    211             } else {
    212                 try self.excess.append(self.aa.allocator(), try self.aa.allocator().dupe(u8, nxt));
    213             }
    214         }
    215     }
    216     for (self.flags.items) |aflag| {
    217         if (aflag.argvalue) |_| {} else if (aflag.defaultargvalue) |dav| {
    218             aflag.argvalue = try self.aa.allocator().dupe(u8, dav);
    219         }
    220     }
    221     for (self.positionals.items) |pos| {
    222         if (pos.value) |_| {} else if (pos.default) |dav| {
    223             pos.value = try self.aa.allocator().dupe(u8, dav);
    224         }
    225     }
    226 }
    227 
    228 pub fn reset(self: *ArgParse) void {
    229     for (self.flags.items) |flag| {
    230         flag.waspresent = false;
    231         flag.argvalue = null;
    232     }
    233     for (self.positionals.items) |pos| {
    234         pos.value = null;
    235     }
    236     for (self.subcommands.items) |sc| {
    237         sc.wasExecuted = false;
    238         for (sc.flags.items) |flag| {
    239             flag.waspresent = false;
    240             flag.argvalue = null;
    241         }
    242         for (sc.positionals.items) |pos| {
    243             pos.value = null;
    244         }
    245     }
    246 }
    247 
    248 const TestArgs = struct {
    249     args: []const []const u8,
    250     pos: usize = 0,
    251     pub fn next(self: *TestArgs) ?[]const u8 {
    252         if (self.pos < self.args.len) {
    253             defer self.pos += 1;
    254             return self.args[self.pos];
    255         } else {
    256             return null;
    257         }
    258     }
    259     pub fn first(self: *TestArgs) []const u8 {
    260         return self.next() orelse @panic("bad test args, no first");
    261     }
    262     pub fn skip(self: *TestArgs) bool {
    263         if (self.next()) |_| {
    264             return true;
    265         } else {
    266             return false;
    267         }
    268     }
    269 };
    270 
    271 test "flag not present" {
    272     const a = std.testing.allocator;
    273     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    274     defer ap.deinit();
    275     var flag: Flag = .{
    276         .long = "myflag",
    277         .short = "m",
    278         .description = "Hello there",
    279         .hasarg = false,
    280     };
    281     try ap.addFlag(&flag);
    282     var ta: TestArgs = .{ .args = &.{ "prog", "foobar" } };
    283     try ap.parseInternal(&ta);
    284     try std.testing.expect(!flag.waspresent);
    285 }
    286 
    287 test "flag with no arg" {
    288     const a = std.testing.allocator;
    289     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    290     defer ap.deinit();
    291     var flag: Flag = .{
    292         .long = "myflag",
    293         .short = "m",
    294         .description = "Hello there",
    295         .hasarg = false,
    296     };
    297     try ap.addFlag(&flag);
    298     var ta: TestArgs = .{ .args = &.{ "prog", "--myflag" } };
    299     try ap.parseInternal(&ta);
    300     try std.testing.expect(flag.waspresent);
    301 }
    302 
    303 test "flag with arg, but arg not supplied" {
    304     const a = std.testing.allocator;
    305     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    306     defer ap.deinit();
    307     var flag: Flag = .{
    308         .long = "myflag",
    309         .short = "m",
    310         .description = "Hello there",
    311         .hasarg = true,
    312     };
    313     try ap.addFlag(&flag);
    314     var ta: TestArgs = .{ .args = &.{ "prog", "--myflag" } };
    315     try std.testing.expectError(error.InvalidArgs, ap.parseInternal(&ta));
    316 }
    317 
    318 test "flag with no arg, but an arg was supplie" {
    319     const a = std.testing.allocator;
    320     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    321     defer ap.deinit();
    322     var flag: Flag = .{
    323         .long = "myflag",
    324         .short = "m",
    325         .description = "Hello there",
    326         .hasarg = false,
    327     };
    328     try ap.addFlag(&flag);
    329     var ta: TestArgs = .{ .args = &.{ "prog", "--myflag=foobar" } };
    330     try std.testing.expectError(error.InvalidArgs, ap.parseInternal(&ta));
    331 }
    332 
    333 test "flag with arg, separated by space" {
    334     const a = std.testing.allocator;
    335     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    336     defer ap.deinit();
    337     var flag: Flag = .{
    338         .long = "myflag",
    339         .short = "m",
    340         .description = "Hello there",
    341         .hasarg = true,
    342     };
    343     try ap.addFlag(&flag);
    344     var ta: TestArgs = .{ .args = &.{ "prog", "--myflag", "myvalue" } };
    345     try ap.parseInternal(&ta);
    346     try std.testing.expect(flag.waspresent);
    347     try std.testing.expectEqualStrings("myvalue", flag.argvalue.?);
    348 }
    349 
    350 test "flag with an arg, separated by an equals" {
    351     const a = std.testing.allocator;
    352     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    353     defer ap.deinit();
    354     var flag: Flag = .{
    355         .long = "myflag",
    356         .short = "m",
    357         .description = "Hello there",
    358         .hasarg = true,
    359     };
    360     try ap.addFlag(&flag);
    361     var ta: TestArgs = .{ .args = &.{ "prog", "--myflag=myvalue" } };
    362     try ap.parseInternal(&ta);
    363     try std.testing.expect(flag.waspresent);
    364     try std.testing.expectEqualStrings("myvalue", flag.argvalue.?);
    365 }
    366 
    367 test "short flag" {
    368     const a = std.testing.allocator;
    369     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    370     defer ap.deinit();
    371     var flag: Flag = .{
    372         .long = "myflag",
    373         .short = "m",
    374         .description = "Hello there",
    375         .hasarg = true,
    376     };
    377     try ap.addFlag(&flag);
    378     var ta: TestArgs = .{ .args = &.{ "prog", "-m=myvalue" } };
    379     try ap.parseInternal(&ta);
    380     try std.testing.expect(flag.waspresent);
    381     try std.testing.expectEqualStrings("myvalue", flag.argvalue.?);
    382 }
    383 
    384 test "unexpected flag" {
    385     const a = std.testing.allocator;
    386     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    387     defer ap.deinit();
    388     var ta: TestArgs = .{ .args = &.{ "prog", "--myflag=myvalue" } };
    389     try std.testing.expectError(error.InvalidArgs, ap.parseInternal(&ta));
    390 }
    391 
    392 test "duplicate flag" {
    393     const a = std.testing.allocator;
    394     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    395     defer ap.deinit();
    396     var flag: Flag = .{
    397         .long = "myflag",
    398         .short = "m",
    399         .description = "Hello there",
    400         .hasarg = true,
    401     };
    402     try ap.addFlag(&flag);
    403     var ta: TestArgs = .{ .args = &.{ "prog", "--myflag=myvalue", "--myflag", "someothervalue" } };
    404     try std.testing.expectError(error.InvalidArgs, ap.parseInternal(&ta));
    405 }
    406 
    407 test "postitional argument not supplied" {
    408     const a = std.testing.allocator;
    409     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    410     defer ap.deinit();
    411     var pos: Positional = .{
    412         .name = "foo",
    413         .description = "A positional argument",
    414     };
    415     try ap.addPositional(&pos);
    416     var ta: TestArgs = .{ .args = &.{"prog"} };
    417     try ap.parseInternal(&ta);
    418     try std.testing.expectEqual(null, pos.value);
    419 }
    420 
    421 test "postitional argument" {
    422     const a = std.testing.allocator;
    423     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    424     defer ap.deinit();
    425     var pos: Positional = .{
    426         .name = "foo",
    427         .description = "A positional argument",
    428     };
    429     try ap.addPositional(&pos);
    430     var ta: TestArgs = .{ .args = &.{ "prog", "bar" } };
    431     try ap.parseInternal(&ta);
    432     try std.testing.expectEqualStrings("bar", pos.value.?);
    433 }
    434 
    435 test "two postitional arguments, first is supplied" {
    436     const a = std.testing.allocator;
    437     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    438     defer ap.deinit();
    439     var pos: Positional = .{
    440         .name = "foo",
    441         .description = "A positional argument",
    442     };
    443     try ap.addPositional(&pos);
    444     var pos2: Positional = .{
    445         .name = "foo2",
    446         .description = "A second positional argument",
    447     };
    448     try ap.addPositional(&pos2);
    449     var ta: TestArgs = .{ .args = &.{ "prog", "bar" } };
    450     try ap.parseInternal(&ta);
    451     try std.testing.expectEqualStrings("bar", pos.value.?);
    452     try std.testing.expectEqual(null, pos2.value);
    453 }
    454 
    455 test "two positional arguments, both are supplied and some extras" {
    456     const a = std.testing.allocator;
    457     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    458     defer ap.deinit();
    459     var pos: Positional = .{
    460         .name = "foo",
    461         .description = "A positional argument",
    462     };
    463     try ap.addPositional(&pos);
    464     var pos2: Positional = .{
    465         .name = "foo2",
    466         .description = "A second positional argument",
    467     };
    468     try ap.addPositional(&pos2);
    469     var ta: TestArgs = .{ .args = &.{ "prog", "bar", "baz", "bing", "bam" } };
    470     try ap.parseInternal(&ta);
    471     try std.testing.expectEqualStrings("bar", pos.value.?);
    472     try std.testing.expectEqualStrings("baz", pos2.value.?);
    473     try std.testing.expectEqual(@as(usize, 2), ap.excess.items.len);
    474     try std.testing.expectEqualStrings("bing", ap.excess.items[0]);
    475     try std.testing.expectEqualStrings("bam", ap.excess.items[1]);
    476 }
    477 
    478 test "positional argument with a default" {
    479     const a = std.testing.allocator;
    480     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    481     defer ap.deinit();
    482     var pos: Positional = .{
    483         .name = "foo",
    484         .description = "A positional argument",
    485     };
    486     try ap.addPositional(&pos);
    487     var pos2: Positional = .{
    488         .name = "foo2",
    489         .description = "A second positional argument",
    490         .default = "defaultbar",
    491     };
    492     try ap.addPositional(&pos2);
    493     var ta: TestArgs = .{ .args = &.{ "prog", "bar" } };
    494     try ap.parseInternal(&ta);
    495     try std.testing.expectEqualStrings("bar", pos.value.?);
    496     try std.testing.expectEqualStrings("defaultbar", pos2.value.?);
    497 }
    498 
    499 test "help" {
    500     const a = std.testing.allocator;
    501     var ap = ArgParse.init(a, "Myprog", "my program to do things");
    502     defer ap.deinit();
    503     var pos1: Positional = .{
    504         .name = "foo",
    505         .description = "A positional argument",
    506     };
    507     try ap.addPositional(&pos1);
    508     var pos2: Positional = .{
    509         .name = "foo2",
    510         .description = "A second positional argument, with a default",
    511         .default = "defaultbar",
    512     };
    513     try ap.addPositional(&pos2);
    514     var flag1: Flag = .{
    515         .long = "myflag",
    516         .short = "m",
    517         .hasarg = true,
    518         .description = "My flag with an argument",
    519     };
    520     try ap.addFlag(&flag1);
    521     var flag2: Flag = .{
    522         .long = "twoflag",
    523         .short = "t",
    524         .hasarg = true,
    525         .description = "My flag with an argument and a default",
    526         .defaultargvalue = "twoflagdefault",
    527     };
    528     try ap.addFlag(&flag2);
    529     var flag3: Flag = .{
    530         .long = "threeflag",
    531         .short = "tf",
    532         .hasarg = false,
    533         .description = "My flag with no argument",
    534     };
    535     try ap.addFlag(&flag3);
    536 
    537     var out = std.ArrayList(u8).init(a);
    538     defer out.deinit();
    539     try ap.help(out.writer());
    540     const expected =
    541         \\Usage: Myprog [foo] [foo2]
    542         \\my program to do things
    543         \\  foo: A positional argument
    544         \\  foo2: A second positional argument, with a default (Default: defaultbar)
    545         \\
    546         \\Flags:
    547         \\  -m,--myflag=ARG  My flag with an argument
    548         \\  -t,--twoflag=ARG My flag with an argument and a default (Default: twoflagdefault)
    549         \\  -tf,--threeflag  My flag with no argument
    550         \\  -h,--help        Show this help
    551         \\
    552     ;
    553     try std.testing.expectEqualStrings(expected, out.items);
    554 }
    555 
    556 test "helpwanted" {
    557     var ap = ArgParse.init(std.testing.allocator, "foo", "foo");
    558     defer ap.deinit();
    559     var ta: TestArgs = .{ .args = &.{ "prog", "--help" } };
    560     try std.testing.expectError(error.HelpWanted, ap.parseInternal(&ta));
    561     var ta2: TestArgs = .{ .args = &.{ "prog", "-h" } };
    562     try std.testing.expectError(error.HelpWanted, ap.parseInternal(&ta2));
    563 }
    564 
    565 test "subcommand" {
    566     var ap = ArgParse.init(std.testing.allocator, "foo", "foo");
    567     defer ap.deinit();
    568     var sc1: Subcommand = .{
    569         .parent = &ap,
    570         .name = "bar",
    571         .description = "Do the bar thing",
    572     };
    573     try ap.addSubcommand(&sc1);
    574     var ta: TestArgs = .{ .args = &.{ "prog", "bar" } };
    575     try ap.parseInternal(&ta);
    576     try std.testing.expect(sc1.wasExecuted);
    577 }
    578 
    579 test "subcommand with flag and positional" {
    580     var ap = ArgParse.init(std.testing.allocator, "foo", "foo");
    581     defer ap.deinit();
    582 
    583     var f1: Flag = .{
    584         .long = "baz",
    585         .hasarg = false,
    586         .description = "baz flag",
    587     };
    588     try ap.addFlag(&f1);
    589     var p1: Positional = .{
    590         .name = "bam",
    591         .description = "do things to bam",
    592     };
    593     try ap.addPositional(&p1);
    594 
    595     var sc1: Subcommand = .{
    596         .parent = &ap,
    597         .name = "bar",
    598         .description = "Do the bar thing",
    599     };
    600     var scf1: Flag = .{
    601         .long = "barflag1",
    602         .hasarg = false,
    603         .description = "barflag1",
    604     };
    605     try sc1.addFlag(&scf1);
    606     var scp1: Positional = .{
    607         .name = "barpos1",
    608         .description = "do bar to this thing",
    609     };
    610     try sc1.addPositional(&scp1);
    611     try ap.addSubcommand(&sc1);
    612 
    613     var ta: TestArgs = .{ .args = &.{"prog"} };
    614     try ap.parseInternal(&ta);
    615     try std.testing.expect(!sc1.wasExecuted);
    616     try std.testing.expect(!f1.waspresent);
    617     try std.testing.expect(p1.value == null);
    618     try std.testing.expect(!scf1.waspresent);
    619     try std.testing.expect(scp1.value == null);
    620 
    621     ap.reset();
    622     ta = .{ .args = &.{ "prog", "--baz", "yoyoyo" } };
    623     try ap.parseInternal(&ta);
    624     try std.testing.expect(!sc1.wasExecuted);
    625     try std.testing.expect(f1.waspresent);
    626     try std.testing.expectEqualStrings("yoyoyo", p1.value.?);
    627     try std.testing.expect(!scf1.waspresent);
    628     try std.testing.expect(scp1.value == null);
    629 
    630     ap.reset();
    631     ta = .{ .args = &.{ "prog", "bar", "yoyoyo", "--barflag1" } };
    632     try ap.parseInternal(&ta);
    633     try std.testing.expect(sc1.wasExecuted);
    634     try std.testing.expect(!f1.waspresent);
    635     try std.testing.expect(p1.value == null);
    636     try std.testing.expect(scf1.waspresent);
    637     try std.testing.expectEqualStrings("yoyoyo", scp1.value.?);
    638 }
    639 
    640 // TODO extend help with subcommands.