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 }