zig/lib/std/progress.zig
Andrew Kelley 0237e7a701 std.io.getStdOut and related fns no longer can error
Thanks to the Windows Process Environment Block, it is possible to
obtain handles to the standard input, output, and error streams without
possibility of failure.
2019-11-13 04:01:40 +00:00

256 lines
9.4 KiB
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();
}
}