refactor: simplify color formatting (#44)

* feat: add ANSI escape code filter writer * refactor(log): use ANSI color filter for printing * refactor: remove redundant Constants.use_color checks * feat: add config option and flag to configure colors
water-sucks · Aug 29, 2024 · 0f1ba04 · 0f1ba04
1 parent 40cab61
commit 0f1ba04
Show file tree

Hide file tree

Showing 11 changed files with 176 additions and 185 deletions.
diff --git a/src/apply.zig b/src/apply.zig
@@ -173,6 +173,10 @@ pub const ApplyCommand = struct {
                 const next_args = try getNextArgs(args, arg, 2);
                 try parsed.build_options.appendSlice(&.{ arg, next_args[0], next_args[1] });
             } else {
+                if (argparse.isFlag(arg)) {
+                    return arg;
+                }
+
                 if (opts.flake and parsed.flake == null) {
                     parsed.flake = arg;
                 } else {

diff --git a/src/config.zig b/src/config.zig
@@ -77,6 +77,7 @@ pub const Config = struct {
         use_nom: bool = false,
         use_git_commit_msg: bool = false,
     } = .{},
+    color: bool = true,
     config_location: []const u8 = "/etc/nixos",
     enter: struct {
         mount_resolv_conf: bool = true,
@@ -206,11 +207,7 @@ pub fn deinit() void {
 }
 
 fn configError(comptime fmt: []const u8, args: anytype) void {
-    if (Constants.use_color) {
-        log.print(ansi.BOLD ++ ansi.RED ++ "error" ++ ansi.RESET ++ ": ", .{});
-    } else {
-        log.print("error: ", .{});
-    }
+    log.print(ansi.BOLD ++ ansi.RED ++ "error" ++ ansi.RESET ++ ": ", .{});
     log.print("invalid setting: ", .{});
     log.print(fmt ++ "\n", args);
 }
diff --git a/src/generation/list.zig b/src/generation/list.zig
@@ -92,7 +92,7 @@ fn listGenerations(allocator: Allocator, profile_name: []const u8, args: Generat
     }
 
     for (generations, 0..) |gen, i| {
-        gen.prettyPrint(.{ .color = Constants.use_color }, stdout) catch unreachable;
+        gen.prettyPrint(.{}, stdout) catch unreachable;
         if (i != generations.len - 1) {
             print(stdout, "\n", .{});
         }

diff --git a/src/info.zig b/src/info.zig
@@ -159,7 +159,6 @@ fn info(allocator: Allocator, args: InfoCommand) InfoError!void {
     }
 
     generation_info.prettyPrint(.{
-        .color = Constants.use_color,
         .show_current_marker = false,
     }, stdout) catch unreachable;
 }

diff --git a/src/log.zig b/src/log.zig
@@ -3,26 +3,32 @@
 //! does not filter output based on build type.
 
 const std = @import("std");
+const io = std.io;
+const mem = std.mem;
 
 const utils = @import("utils.zig");
 const ansi = utils.ansi;
+const ANSIFilter = ansi.ANSIFilter;
 
-const Constants = @import("constants.zig");
+/// Print to stderr. This makes sure that ANSI codes are handled
+/// according to whether or not they are disabled.
+pub fn print(comptime fmt: []const u8, args: anytype) void {
+    const stderr = std.io.getStdErr().writer();
+    std.debug.lockStdErr();
+    defer std.debug.unlockStdErr();
+
+    var color_filter = ANSIFilter(@TypeOf(stderr)){ .raw_writer = stderr };
+    const writer = color_filter.writer();
+    writer.print(fmt, args) catch return;
+}
 
 /// Base logging function with no level. Prints a newline automatically.
 fn log(comptime prefix: []const u8, comptime fmt: []const u8, args: anytype) void {
     const real_prefix = prefix ++ (if (prefix.len != 0) ": " else "");
 
-    const stderr = std.io.getStdErr().writer();
-    std.debug.lockStdErr();
-    defer std.debug.unlockStdErr();
-    nosuspend stderr.print(real_prefix ++ fmt ++ "\n", args) catch return;
+    print(real_prefix ++ fmt ++ "\n", args);
 }
 
-/// Bare print that gets rid of the `std.debug` prefix,
-/// which is clunky.
-pub const print = std.debug.print;
-
 /// Global step counter. This will increase every single time step() is invoked.
 var step_num: usize = 0;
 
@@ -32,57 +38,36 @@ pub fn step(comptime fmt: []const u8, args: anytype) void {
     if (step_num > 1) {
         print("\n", .{});
     }
-    if (Constants.use_color) {
-        print(ansi.BOLD ++ ansi.MAGENTA, .{});
-    }
+
+    print(ansi.BOLD ++ ansi.MAGENTA, .{});
 
     print("{d}. ", .{step_num});
     print(fmt, args);
 
-    if (Constants.use_color) {
-        print(ansi.RESET, .{});
-    }
-    print("\n", .{});
+    print(ansi.RESET ++ "\n", .{});
 }
 
 /// Pretty-print a command that will be ran.
 pub fn cmd(argv: []const []const u8) void {
-    if (Constants.use_color) {
-        print(ansi.BLUE, .{});
-    }
+    print(ansi.BLUE, .{});
 
-    print("$ ", .{});
+    print(ansi.BLUE ++ "$ ", .{});
     for (argv) |arg| {
         print("{s} ", .{arg});
     }
-
-    if (Constants.use_color) {
-        print(ansi.RESET, .{});
-    }
+    print(ansi.RESET, .{});
 
     print("\n", .{});
 }
 
 pub fn err(comptime fmt: []const u8, args: anytype) void {
-    if (Constants.use_color) {
-        log(ansi.BOLD ++ ansi.RED ++ "error" ++ ansi.RESET, fmt, args);
-    } else {
-        log("error", fmt, args);
-    }
+    log(ansi.BOLD ++ ansi.RED ++ "error" ++ ansi.RESET, fmt, args);
 }
 
 pub fn warn(comptime fmt: []const u8, args: anytype) void {
-    if (Constants.use_color) {
-        log(ansi.BOLD ++ ansi.YELLOW ++ "warning" ++ ansi.RESET, fmt, args);
-    } else {
-        log("warning", fmt, args);
-    }
+    log(ansi.BOLD ++ ansi.YELLOW ++ "warning" ++ ansi.RESET, fmt, args);
 }
 
 pub fn info(comptime fmt: []const u8, args: anytype) void {
-    if (Constants.use_color) {
-        log(ansi.GREEN ++ "info" ++ ansi.RESET, fmt, args);
-    } else {
-        log("info", fmt, args);
-    }
+    log(ansi.GREEN ++ "info" ++ ansi.RESET, fmt, args);
 }
diff --git a/src/main.zig b/src/main.zig
@@ -51,6 +51,7 @@ const println = utils.println;
 const MainArgs = struct {
     allocator: Allocator,
     config_values: ArrayList([]const u8),
+    color_always: bool = false,
     subcommand: ?Subcommand = null,
 
     const Self = @This();
@@ -91,6 +92,7 @@ const MainArgs = struct {
         \\
         \\Options:
         \\    -c, --config <KEY>=<VALUE>    Set a configuration value
+        \\    -C, --color-always            Always color output when possible
         \\    -h, --help                    Show this help menu
         \\    -v, --version                 Print version information
         \\
@@ -112,6 +114,8 @@ const MainArgs = struct {
             if (argparse.argIs(arg, "--config", "-c")) {
                 const next = (try getNextArgs(argv, arg, 1))[0];
                 try result.config_values.append(next);
+            } else if (argparse.argIs(arg, "--color-always", "-C")) {
+                result.color_always = true;
             } else if (argparse.argIs(arg, "--help", "-h")) {
                 log.print(usage, .{});
                 return ArgParseError.HelpInvoked;
@@ -239,7 +243,10 @@ pub fn main() !u8 {
     config.parseConfig(allocator) catch {};
     defer config.deinit();
 
-    Constants.use_color = !mem.eql(u8, posix.getenv("NO_COLOR") orelse "", "1") and
+    // This initially sets NO_COLOR to a reasonable value before the configuration
+    // and arguments are parsed. This will be re-set later.
+    const no_color_is_set = mem.eql(u8, posix.getenv("NO_COLOR") orelse "", "1");
+    Constants.use_color = !no_color_is_set and
         io.getStdOut().supportsAnsiEscapeCodes() and
         io.getStdErr().supportsAnsiEscapeCodes();
 
@@ -277,6 +284,24 @@ pub fn main() !u8 {
     }
     config.validateConfig(allocator) catch return 2;
 
+    const c = config.getConfig();
+
+    // Now that we have the real color settings from parsing
+    // the configuration and command-line arguments, set it.
+    //
+    // Precedence of color settings:
+    // 1. -C flag -> true
+    // 2. NO_COLOR=1 -> false
+    // 3. `color` setting from config (default: true)
+    if (structured_args.color_always) {
+        Constants.use_color = true;
+    } else if (no_color_is_set) {
+        Constants.use_color = false;
+    } else {
+        Constants.use_color = c.color and io.getStdOut().supportsAnsiEscapeCodes() and
+            io.getStdErr().supportsAnsiEscapeCodes();
+    }
+
     const status = switch (structured_args.subcommand.?) {
         .aliases => |args| {
             alias.printAliases(args);

diff --git a/src/option.zig b/src/option.zig
@@ -286,10 +286,7 @@ fn displayOption(allocator: Allocator, opt: NixosOption, evaluated: EvaluatedVal
             break :blk mem.trim(u8, rendered, "\n ");
         }
 
-        break :blk if (Constants.use_color)
-            ansi.ITALIC ++ "(none)" ++ ansi.RESET
-        else
-            "(none)";
+        break :blk ansi.ITALIC ++ "(none)" ++ ansi.RESET;
     };
     defer if (desc_alloc) allocator.free(description);
 
@@ -301,61 +298,36 @@ fn displayOption(allocator: Allocator, opt: NixosOption, evaluated: EvaluatedVal
     };
     const example = if (opt.example) |e| mem.trim(u8, e.text, "\n ") else null;
 
-    if (Constants.use_color) {
-        println(stdout, ansi.BOLD ++ "Name\n" ++ ansi.RESET ++ "{s}\n", .{opt.name});
-        println(stdout, ansi.BOLD ++ "Description\n" ++ ansi.RESET ++ "{s}\n", .{description});
-        println(stdout, ansi.BOLD ++ "Type\n" ++ ansi.RESET ++ ansi.ITALIC ++ "{s}\n" ++ ansi.RESET, .{opt.type});
-        println(stdout, ansi.BOLD ++ "Value" ++ ansi.RESET, .{});
-        if (std.meta.activeTag(evaluated) == .success) {
-            println(stdout, "{s}\n", .{evaluated.success});
-        } else {
-            println(stdout, ansi.RED ++ "error: {s}\n" ++ ansi.RESET, .{evaluated.@"error"});
-        }
-
-        println(stdout, ansi.BOLD ++ "Default" ++ ansi.RESET, .{});
-        if (opt.default) |_| {
-            println(stdout, ansi.WHITE ++ "{s}" ++ ansi.RESET, .{default});
-        } else {
-            println(stdout, ansi.ITALIC ++ "(none)" ++ ansi.RESET, .{});
-        }
-        println(stdout, "", .{});
+    println(stdout, ansi.BOLD ++ "Name\n" ++ ansi.RESET ++ "{s}\n", .{opt.name});
+    println(stdout, ansi.BOLD ++ "Description\n" ++ ansi.RESET ++ "{s}\n", .{description});
+    println(stdout, ansi.BOLD ++ "Type\n" ++ ansi.RESET ++ ansi.ITALIC ++ "{s}\n" ++ ansi.RESET, .{opt.type});
+    println(stdout, ansi.BOLD ++ "Value" ++ ansi.RESET, .{});
+    if (std.meta.activeTag(evaluated) == .success) {
+        println(stdout, "{s}\n", .{evaluated.success});
+    } else {
+        println(stdout, ansi.RED ++ "error: {s}\n" ++ ansi.RESET, .{evaluated.@"error"});
+    }
 
-        if (example) |e| {
-            println(stdout, ansi.BOLD ++ "Example\n" ++ ansi.RESET ++ "{s}\n", .{e});
-        }
-        if (opt.declarations.len > 0) {
-            println(stdout, ansi.BOLD ++ "Declared In" ++ ansi.RESET, .{});
-            for (opt.declarations) |decl| {
-                println(stdout, ansi.ITALIC ++ "  - {s}" ++ ansi.RESET, .{decl});
-            }
-        }
-        if (opt.readOnly) {
-            println(stdout, ansi.YELLOW ++ "\nThis option is read-only." ++ ansi.RESET, .{});
-        }
+    println(stdout, ansi.BOLD ++ "Default" ++ ansi.RESET, .{});
+    if (opt.default) |_| {
+        println(stdout, ansi.WHITE ++ "{s}" ++ ansi.RESET, .{default});
     } else {
-        println(stdout, "Name\n{s}\n", .{opt.name});
-        println(stdout, "Description\n{s}\n", .{description});
-        println(stdout, "Type\n{s}\n", .{opt.type});
-        println(stdout, "Value", .{});
-        if (std.meta.activeTag(evaluated) == .success) {
-            println(stdout, "{s}\n", .{evaluated.success});
-        } else {
-            println(stdout, "error: {s}\n", .{evaluated.@"error"});
-        }
-        println(stdout, "Default\n{s}\n", .{default});
-        if (example) |e| {
-            println(stdout, "Example\n{s}\n", .{e});
-        }
-        if (opt.declarations.len > 0) {
-            println(stdout, "Declared In", .{});
-            for (opt.declarations) |decl| {
-                println(stdout, "  - {s}", .{decl});
-            }
-        }
-        if (opt.readOnly) {
-            println(stdout, "\nThis option is read-only.", .{});
+        println(stdout, ansi.ITALIC ++ "(none)" ++ ansi.RESET, .{});
+    }
+    println(stdout, "", .{});
+
+    if (example) |e| {
+        println(stdout, ansi.BOLD ++ "Example\n" ++ ansi.RESET ++ "{s}\n", .{e});
+    }
+    if (opt.declarations.len > 0) {
+        println(stdout, ansi.BOLD ++ "Declared In" ++ ansi.RESET, .{});
+        for (opt.declarations) |decl| {
+            println(stdout, ansi.ITALIC ++ "  - {s}" ++ ansi.RESET, .{decl});
         }
     }
+    if (opt.readOnly) {
+        println(stdout, ansi.YELLOW ++ "\nThis option is read-only." ++ ansi.RESET, .{});
+    }
 }
 
 pub fn evaluateOptionValue(allocator: Allocator, configuration: ConfigType, name: []const u8) !EvaluatedValue {

diff --git a/src/repl.zig b/src/repl.zig
@@ -19,8 +19,6 @@ const argError = argparse.argError;
 const argIs = argparse.argIs;
 const getNextArgs = argparse.getNextArgs;
 
-const Constants = @import("constants.zig");
-
 const log = @import("log.zig");
 
 const utils = @import("utils.zig");
@@ -133,7 +131,7 @@ const flake_motd_template =
 
 fn execFlakeRepl(allocator: Allocator, ref: FlakeRef, includes: []const []const u8) !void {
     // zig fmt: off
-    const motd = if (Constants.use_color) try fmt.allocPrint(allocator, flake_motd_template, .{
+    const motd = try fmt.allocPrint(allocator, flake_motd_template, .{
         ansi.CYAN,    ref.uri,    ref.system,   ansi.RESET,
         ansi.MAGENTA, ansi.RESET,
         ansi.MAGENTA, ansi.RESET,
@@ -143,16 +141,6 @@ fn execFlakeRepl(allocator: Allocator, ref: FlakeRef, includes: []const []const
         ansi.GREEN,   ansi.RESET,
         ansi.GREEN,   ansi.RESET,
         ansi.YELLOW,  ansi.RESET, ansi.CYAN, ansi.RESET,
-    }) else try fmt.allocPrint(allocator, flake_motd_template, .{
-        "", ref.uri, ref.system, "",
-        "",      "",
-        "",      "",
-        "",      "",
-        "",      "",         "", "",
-        "",      "",         "", "",
-        "",      "",
-        "",      "",
-        "",      "",         "", ""
     });
     // zig fmt: on
 
@@ -203,25 +191,13 @@ fn execLegacyRepl(allocator: Allocator, includes: []const []const u8, impure: bo
     defer argv.deinit();
 
     // zig fmt: off
-    const motd = if (Constants.use_color) try fmt.allocPrint(allocator, legacy_motd_template, .{
+    const motd = try fmt.allocPrint(allocator, legacy_motd_template, .{
         ansi.MAGENTA, ansi.RESET,
         ansi.MAGENTA, ansi.RESET,
         ansi.MAGENTA, ansi.RESET, ansi.CYAN,    ansi.RESET,
         ansi.MAGENTA, ansi.RESET, ansi.MAGENTA, ansi.RESET,
         ansi.GREEN,   ansi.RESET,
         ansi.GREEN,   ansi.RESET,
-    }) else try fmt.allocPrint(allocator, flake_motd_template, .{
-        "", "",
-        "", "",
-        "", "", "", "",
-        "", "", "", "",
-        "", "",
-        "", "",
-        "", "",
-        "", "",
-        "", "",
-        "", "",
-        "", "",
     });
     // zig fmt: on