From da9b1159a7257c70af54a361db81b1488ab9113b Mon Sep 17 00:00:00 2001
From: Rich FitzJohn <r.fitzjohn@imperial.ac.uk>
Date: Wed, 6 Nov 2024 07:51:39 +0000
Subject: [PATCH] Printing improvements

---
 R/progress.R                  |  2 +-
 R/samples.R                   |  4 ++++
 tests/testthat/test-samples.R | 40 +++++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/R/progress.R b/R/progress.R
index 5577db75..3ebc567f 100644
--- a/R/progress.R
+++ b/R/progress.R
@@ -62,7 +62,7 @@ progress_bar_fancy <- function(n_chains, n_steps, show_overall,
                "{cli::pb_percent} ETA: {cli::pb_eta}")
   fmt_done <- paste(
     "{cli::col_green(cli::symbol$tick)} Sampled {cli::pb_total} steps",
-    "across {n_chains} chains in {cli::pb_elapsed}")
+    "across {n_chains} chain{?s} in {cli::pb_elapsed}")
   fmt_failed <- paste(
     "{cli::col_red(cli::symbol$cross)} Sampling stopped at {cli::pb_current}",
     "step{?s} after {cli::pb_elapsed}")
diff --git a/R/samples.R b/R/samples.R
index 80e675ff..691e3aa0 100644
--- a/R/samples.R
+++ b/R/samples.R
@@ -26,6 +26,10 @@ print.monty_samples <- function(x, ...) {
     cli::cli_alert_info(
       "These samples can be restared with {.help monty_sample_continue}")
   }
+  if (!is.null(x$observations)) {
+    cli::cli_alert_info(
+      "These samples have associated observations")
+  }
 
   cli::cli_alert_info(
     paste('See {.help monty_sample} and {.run vignette("samples")} for more',
diff --git a/tests/testthat/test-samples.R b/tests/testthat/test-samples.R
index 85943021..d7e327e3 100644
--- a/tests/testthat/test-samples.R
+++ b/tests/testthat/test-samples.R
@@ -12,6 +12,46 @@ test_that("can print samples", {
 })
 
 
+test_that("print information about observations", {
+  model <- ex_simple_gamma1()
+  sampler <- monty_sampler_random_walk(vcv = diag(1) * 0.01)
+  s <- monty_sample(model, sampler, 100, 1, n_chains = 3)
+
+  res <- evaluate_promise(print(s))
+  expect_no_match(
+    res$messages,
+    "These samples have associated observations",
+    fixed = TRUE, all = FALSE)
+
+  s$observations <- list(a = 1)
+  res <- evaluate_promise(print(s))
+  expect_match(
+    res$messages,
+    "These samples have associated observations",
+    fixed = TRUE, all = FALSE)
+})
+
+
+test_that("print information about restart", {
+  model <- ex_simple_gamma1()
+  sampler <- monty_sampler_random_walk(vcv = diag(1) * 0.01)
+  s <- monty_sample(model, sampler, 100, 1, n_chains = 3)
+
+  res <- evaluate_promise(print(s))
+  expect_no_match(
+    res$messages,
+    "These samples can be restared with",
+    fixed = TRUE, all = FALSE)
+
+  s$restart <- list()
+  res <- evaluate_promise(print(s))
+  expect_match(
+    res$messages,
+    "These samples can be restared with",
+    fixed = TRUE, all = FALSE)
+})
+
+
 test_that("can convert to posterior draws types", {
   skip_if_not_installed("posterior")
   model <- ex_simple_gamma1()