zigvm

zigvm: Zig Version Manager
git clone git://code.mfashby.net:/zigvm
Log | Files | Refs | README

argparse.zig (16469B)


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