zig/lib/std/progress.zig

const std = @import("std");
const testing = std.testing;
const assert = std.debug.assert;

/// This API is non-allocating and non-fallible. The tradeoff is that users of
/// this API must provide the storage for each `Progress.Node`.
/// Initialize the struct directly, overriding these fields as desired:
/// * `refresh_rate_ms`
/// * `initial_delay_ms`
pub const Progress = struct {
    /// `null` if the current node (and its children) should
    /// not print on update()
    terminal: ?std.fs.File = undefined,

    root: Node = undefined,

    /// Keeps track of how much time has passed since the beginning.
    /// Used to compare with `initial_delay_ms` and `refresh_rate_ms`.
    timer: std.time.Timer = undefined,

    /// When the previous refresh was written to the terminal.
    /// Used to compare with `refresh_rate_ms`.
    prev_refresh_timestamp: u64 = undefined,

    /// This buffer represents the maximum number of bytes written to the terminal
    /// with each refresh.
    output_buffer: [100]u8 = undefined,

    /// How many nanoseconds between writing updates to the terminal.
    refresh_rate_ns: u64 = 50 * std.time.millisecond,

    /// How many nanoseconds to keep the output hidden
    initial_delay_ns: u64 = 500 * std.time.millisecond,

    done: bool = true,

    /// Keeps track of how many columns in the terminal have been output, so that
    /// we can move the cursor back later.
    columns_written: usize = undefined,

    /// Represents one unit of progress. Each node can have children nodes, or
    /// one can use integers with `update`.
    pub const Node = struct {
        context: *Progress,
        parent: ?*Node,
        completed_items: usize,
        name: []const u8,
        recently_updated_child: ?*Node = null,

        /// This field may be updated freely.
        estimated_total_items: ?usize,

        /// Create a new child progress node.
        /// Call `Node.end` when done.
        /// TODO solve https://github.com/ziglang/zig/issues/2765 and then change this
        /// API to set `self.parent.recently_updated_child` with the return value.
        /// Until that is fixed you probably want to call `activate` on the return value.
        pub fn start(self: *Node, name: []const u8, estimated_total_items: ?usize) Node {
            return Node{
                .context = self.context,
                .parent = self,
                .completed_items = 0,
                .name = name,
                .estimated_total_items = estimated_total_items,
            };
        }

        /// This is the same as calling `start` and then `end` on the returned `Node`.
        pub fn completeOne(self: *Node) void {
            if (self.parent) |parent| parent.recently_updated_child = self;
            self.completed_items += 1;
            self.context.maybeRefresh();
        }

        pub fn end(self: *Node) void {
            self.context.maybeRefresh();
            if (self.parent) |parent| {
                if (parent.recently_updated_child) |parent_child| {
                    if (parent_child == self) {
                        parent.recently_updated_child = null;
                    }
                }
                parent.completeOne();
            } else {
                self.context.done = true;
                self.context.refresh();
            }
        }

        /// Tell the parent node that this node is actively being worked on.
        pub fn activate(self: *Node) void {
            if (self.parent) |parent| parent.recently_updated_child = self;
        }
    };

    /// Create a new progress node.
    /// Call `Node.end` when done.
    /// TODO solve https://github.com/ziglang/zig/issues/2765 and then change this
    /// API to return Progress rather than accept it as a parameter.
    pub fn start(self: *Progress, name: []const u8, estimated_total_items: ?usize) !*Node {
        const stderr = std.io.getStdErr();
        self.terminal = if (stderr.supportsAnsiEscapeCodes()) stderr else null;
        self.root = Node{
            .context = self,
            .parent = null,
            .completed_items = 0,
            .name = name,
            .estimated_total_items = estimated_total_items,
        };
        self.columns_written = 0;
        self.prev_refresh_timestamp = 0;
        self.timer = try std.time.Timer.start();
        self.done = false;
        return &self.root;
    }

    /// Updates the terminal if enough time has passed since last update.
    pub fn maybeRefresh(self: *Progress) void {
        const now = self.timer.read();
        if (now < self.initial_delay_ns) return;
        if (now - self.prev_refresh_timestamp < self.refresh_rate_ns) return;
        self.refresh();
    }

    /// Updates the terminal and resets `self.next_refresh_timestamp`.
    pub fn refresh(self: *Progress) void {
        const file = self.terminal orelse return;

        const prev_columns_written = self.columns_written;
        var end: usize = 0;
        if (self.columns_written > 0) {
            // restore cursor position
            end += (std.fmt.bufPrint(self.output_buffer[end..], "\x1b[{}D", self.columns_written) catch unreachable).len;
            self.columns_written = 0;

            // clear rest of line
            end += (std.fmt.bufPrint(self.output_buffer[end..], "\x1b[0K") catch unreachable).len;
        }

        if (!self.done) {
            var need_ellipse = false;
            var maybe_node: ?*Node = &self.root;
            while (maybe_node) |node| {
                if (need_ellipse) {
                    self.bufWrite(&end, "...");
                }
                need_ellipse = false;
                if (node.name.len != 0 or node.estimated_total_items != null) {
                    if (node.name.len != 0) {
                        self.bufWrite(&end, "{}", node.name);
                        need_ellipse = true;
                    }
                    if (node.estimated_total_items) |total| {
                        if (need_ellipse) self.bufWrite(&end, " ");
                        self.bufWrite(&end, "[{}/{}] ", node.completed_items + 1, total);
                        need_ellipse = false;
                    } else if (node.completed_items != 0) {
                        if (need_ellipse) self.bufWrite(&end, " ");
                        self.bufWrite(&end, "[{}] ", node.completed_items + 1);
                        need_ellipse = false;
                    }
                }
                maybe_node = node.recently_updated_child;
            }
            if (need_ellipse) {
                self.bufWrite(&end, "...");
            }
        }

        _ = file.write(self.output_buffer[0..end]) catch |e| {
            // Stop trying to write to this file once it errors.
            self.terminal = null;
        };
        self.prev_refresh_timestamp = self.timer.read();
    }

    pub fn log(self: *Progress, comptime format: []const u8, args: ...) void {
        const file = self.terminal orelse return;
        self.refresh();
        file.outStream().stream.print(format, args) catch {
            self.terminal = null;
            return;
        };
        self.columns_written = 0;
    }

    fn bufWrite(self: *Progress, end: *usize, comptime format: []const u8, args: ...) void {
        if (std.fmt.bufPrint(self.output_buffer[end.*..], format, args)) |written| {
            const amt = written.len;
            end.* += amt;
            self.columns_written += amt;
        } else |err| switch (err) {
            error.BufferTooSmall => {
                self.columns_written += self.output_buffer.len - end.*;
                end.* = self.output_buffer.len;
            },
        }
        const bytes_needed_for_esc_codes_at_end = 11;
        const max_end = self.output_buffer.len - bytes_needed_for_esc_codes_at_end;
        if (end.* > max_end) {
            const suffix = "...";
            self.columns_written = self.columns_written - (end.* - max_end) + suffix.len;
            std.mem.copy(u8, self.output_buffer[max_end..], suffix);
            end.* = max_end + suffix.len;
        }
    }
};

test "basic functionality" {
    var disable = true;
    if (disable) {
        // This test is disabled because it uses time.sleep() and is therefore slow. It also
        // prints bogus progress data to stderr.
        return error.SkipZigTest;
    }
    var progress = Progress{};
    const root_node = try progress.start("", 100);
    defer root_node.end();

    const sub_task_names = [_][]const u8{
        "reticulating splines",
        "adjusting shoes",
        "climbing towers",
        "pouring juice",
    };
    var next_sub_task: usize = 0;

    var i: usize = 0;
    while (i < 100) : (i += 1) {
        var node = root_node.start(sub_task_names[next_sub_task], 5);
        node.activate();
        next_sub_task = (next_sub_task + 1) % sub_task_names.len;

        node.completeOne();
        std.time.sleep(5 * std.time.millisecond);
        node.completeOne();
        node.completeOne();
        std.time.sleep(5 * std.time.millisecond);
        node.completeOne();
        node.completeOne();
        std.time.sleep(5 * std.time.millisecond);

        node.end();

        std.time.sleep(5 * std.time.millisecond);
    }
    {
        var node = root_node.start("this is a really long name designed to activate the truncation code. let's find out if it works", null);
        node.activate();
        std.time.sleep(10 * std.time.millisecond);
        progress.refresh();
        std.time.sleep(10 * std.time.millisecond);
        node.end();
    }
}