diff --git a/barkeep/barkeep.h b/barkeep/barkeep.h index 221db7d..75c838b 100644 --- a/barkeep/barkeep.h +++ b/barkeep/barkeep.h @@ -388,6 +388,12 @@ class Status : public Animation { std::lock_guard lock(message_mutex_); message_ = message; } + + /// Get the current message. + std::string message() { + std::lock_guard lock(message_mutex_); + return message_; + } }; /// Creates a composite display out of two display that shows them side by side. diff --git a/python/barkeep.cpp b/python/barkeep.cpp index eaf58c1..934eaab 100644 --- a/python/barkeep.cpp +++ b/python/barkeep.cpp @@ -82,6 +82,42 @@ class Animation_ : public Animation { } }; +class Status_ : public Status { + public: + std::shared_ptr file_ = nullptr; + + Status_(py::object file = py::none(), + std::string message = "", + std::variant style = Ellipsis, + double interval = 0., + bool no_tty = false) + : Status({.out = nullptr, + .message = message, + .style = style, + .interval = interval, + .no_tty = no_tty, + .show = false}) { + if (not file.is_none()) { + file_ = std::make_shared(std::move(file)); + } + out_ = file_ ? (std::ostream*)file_.get() : &std::cout; + } + + void join() override { + if (file_) { + // release gil because displayer thread needs it to write + py::gil_scoped_release release; + AsyncDisplay::join(); + } else { + AsyncDisplay::join(); + } + } + + std::unique_ptr clone() const override { + return std::make_unique(*this); + } +}; + template class Counter_ : public Counter { protected: @@ -303,7 +339,8 @@ PYBIND11_MODULE(barkeep, m) { std::variant style, bool no_tty, bool show) { - auto a = std::make_unique(file, msg, style, interval, no_tty); + auto a = std::make_unique( + file, msg, style, interval, no_tty); if (show) { a->show(); } return a; }), @@ -333,7 +370,51 @@ PYBIND11_MODULE(barkeep, m) { "no_tty"_a = false, "show"_a = true, py::keep_alive<0, 1>()); // keep file alive while the animation is - // alive); + // alive + + py::class_(m, "Status") + .def(py::init([](py::object file, + std::string msg, + double interval, + std::variant style, + bool no_tty, + bool show) { + auto a = + std::make_unique(file, msg, style, interval, no_tty); + if (show) { a->show(); } + return a; + }), + R"docstr( + Status is an Animation where it is possible to update the message + while the animation is running. + + Parameters + ---------- + file : file-like object, optional + File to write to. Defaults to stdout. + message : str, optional + Message to display. Defaults to "". + interval : float, optional + Interval between frames in seconds. If None, defaults to 1 if + not no_tty, 60 otherwise. + style : AnimationStyle, optional + Animation style. Defaults to AnimationStyle.Ellipsis. + no_tty : bool, optional + If True, use no-tty mode (no \r, slower refresh). Defaults to False. + show : bool, optional + If True, show the animation immediately. Defaults to True. + )docstr", + "file"_a = py::none(), + "message"_a = "", + "interval"_a = 1., + "style"_a = AnimationStyle::Ellipsis, + "no_tty"_a = false, + "show"_a = true, + py::keep_alive<0, 1>()) // keep file alive while the animation is + // alive + .def_property("message", + py::overload_cast<>(&Status_::message), + py::overload_cast(&Status_::message)); auto bind_display = [&](auto& m, auto disp, auto pv, const char* name) { using T = decltype(pv); @@ -558,7 +639,7 @@ PYBIND11_MODULE(barkeep, m) { if (self.running() or other.running()) { // not sure why this is necessary, but it prevents segfaults. // maybe pybind11 implicit copies are causing problems when destructor - // attempts a done() ? + // attempts a `done()`? self.done(); other.done(); throw std::runtime_error("Cannot combine running AsyncDisplay objects!"); diff --git a/python/tests/demo.py b/python/tests/demo.py index 48bc6ed..976add2 100644 --- a/python/tests/demo.py +++ b/python/tests/demo.py @@ -5,6 +5,7 @@ DType, ProgressBar, ProgressBarStyle, + Status, ) from argparse import ArgumentParser from collections import OrderedDict @@ -133,9 +134,21 @@ def fmt(): work.done() +def status(): + s = Status(message="Working") + time.sleep(2.5) + s.message = "Still working" + time.sleep(2.5) + s.message = "Almost done" + time.sleep(2.5) + s.message = "Done" + s.done() + + demos = OrderedDict( [ ("animation", animation), + ("status", status), ("counter", counter), ("progress_bar", progress_bar), ("composite", composite), @@ -148,7 +161,15 @@ def fmt(): parser = ArgumentParser() parser.add_argument( "demo", - choices=["animation", "counter", "progress_bar", "composite", "fmt", []], + choices=[ + "animation", + "status", + "counter", + "progress_bar", + "composite", + "fmt", + [], + ], nargs="*", ) diff --git a/python/tests/test.py b/python/tests/test.py index 078867c..6f350c6 100644 --- a/python/tests/test.py +++ b/python/tests/test.py @@ -6,6 +6,7 @@ DType, ProgressBar, ProgressBarStyle, + Status, ) import pytest import random @@ -49,6 +50,18 @@ def check_anim(parts: list[str], msg: str, stills: list[str]): assert part == (msg + " " + stills[j] + " ") +def check_status(parts: list[str], messages: list[str], stills: list[str]): + msg_i = 0 + for i in range(len(parts) - 1): + j = i % len(stills) + part = parts[i] + msg = messages[msg_i] + if part != (msg + " " + stills[j] + " "): + msg_i += 1 + msg = messages[msg_i] + assert part == (msg + " " + stills[j] + " ") + + animation_styles = [ AnimationStyle.Ellipsis, AnimationStyle.Clock, @@ -90,6 +103,24 @@ def test_animation(i: int, sty: AnimationStyle): check_anim(check_and_get_parts(out.getvalue()), "Working", animation_stills[i]) +@pytest.mark.parametrize("i,sty", enumerate(animation_styles)) +def test_status(i: int, sty: AnimationStyle): + out = io.StringIO() + + stat = Status(message="Working", style=sty, interval=0.1, file=out) + time.sleep(0.5) + stat.message = "Still working" + time.sleep(0.5) + stat.message = "Done" + stat.done() + + check_status( + check_and_get_parts(out.getvalue()), + ["Working", "Still working", "Done"], + animation_stills[i], + ) + + def test_custom_animation(): out = io.StringIO() diff --git a/tests/test.cpp b/tests/test.cpp index e1d8eb7..5e841f9 100644 --- a/tests/test.cpp +++ b/tests/test.cpp @@ -157,6 +157,7 @@ TEST_CASE("Status", "[status]") { check_status(parts, {"Working", "Still working", "Done"}, animation_stills_[size_t(sty)].first); + CHECK(stat.message() == "Done"); } using ProgressTypeList =