From f5bfbccfc84c5d207a1c5e06f66adda3d4fce1e3 Mon Sep 17 00:00:00 2001
From: Maxim Kolesnikov <indev.src@gmail.com>
Date: Wed, 23 Feb 2022 21:01:28 +0700
Subject: [PATCH] feat(script): add formatting for script failure (#2596)

Closes #2588

* feat(script): add formatting for script failure

* refactor: address review comments

* doc: add changelog entry

* refactor: minor cleanup
---
 CHANGELOG.md                       |  1 +
 include/adapters/script_runner.hpp |  4 +++-
 include/modules/script.hpp         |  7 ++++++-
 src/adapters/script_runner.cpp     | 14 +++++++++-----
 src/modules/script.cpp             | 21 ++++++++++++++++++++-
 5 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4664eb3e..959ba70f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -81,6 +81,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `DEBUG_SHADED` cmake variable and its associated functionality.
 
 ### Added
+- `custom/script`: Add formatting for script failure (`format-fail`, `label-fail`) ([`#2588`](https://github.com/polybar/polybar/issues/2588))
 - Support `px` and `pt` units everyhwere where before only a number of spaces
   or pixels could be specified.
   ([`#2578`](https://github.com/polybar/polybar/pull/2578))
diff --git a/include/adapters/script_runner.hpp b/include/adapters/script_runner.hpp
index 9ecec1c1..a2322e50 100644
--- a/include/adapters/script_runner.hpp
+++ b/include/adapters/script_runner.hpp
@@ -24,13 +24,14 @@ class script_runner {
 
   int get_pid() const;
   int get_counter() const;
+  int get_exit_status() const;
 
   string get_output();
 
   bool is_stopping() const;
 
  protected:
-  bool set_output(const string&&);
+  bool set_output(string&&);
 
   interval run_tail();
   interval run();
@@ -52,6 +53,7 @@ class script_runner {
   std::atomic_int m_counter{0};
   std::atomic_bool m_stopping{false};
   std::atomic_int m_pid{-1};
+  std::atomic_int m_exit_status{0};
 };
 
 POLYBAR_NS_END
diff --git a/include/modules/script.hpp b/include/modules/script.hpp
index c00dcdbf..94308249 100644
--- a/include/modules/script.hpp
+++ b/include/modules/script.hpp
@@ -16,6 +16,8 @@ namespace modules {
     void stop() override;
 
     string get_output();
+    string get_format() const;
+
     bool build(builder* builder, const string& tag) const;
 
     static constexpr auto TYPE = "custom/script";
@@ -24,7 +26,9 @@ namespace modules {
     bool check_condition();
 
    private:
-    static constexpr const char* TAG_LABEL{"<label>"};
+    static constexpr auto TAG_LABEL = "<label>";
+    static constexpr auto TAG_LABEL_FAIL = "<label-fail>";
+    static constexpr auto FORMAT_FAIL = "format-fail";
 
     const bool m_tail;
     const script_runner::interval m_interval{0};
@@ -34,6 +38,7 @@ namespace modules {
     map<mousebtn, string> m_actions;
 
     label_t m_label;
+    label_t m_label_fail;
   };
 }  // namespace modules
 
diff --git a/src/adapters/script_runner.cpp b/src/adapters/script_runner.cpp
index e922c263..ac8dd041 100644
--- a/src/adapters/script_runner.cpp
+++ b/src/adapters/script_runner.cpp
@@ -60,6 +60,10 @@ int script_runner::get_counter() const {
   return m_counter;
 }
 
+int script_runner::get_exit_status() const {
+  return m_exit_status;
+}
+
 string script_runner::get_output() {
   std::lock_guard<std::mutex> guard(m_output_lock);
   return m_output;
@@ -74,7 +78,7 @@ bool script_runner::is_stopping() const {
  *
  * Returns true if the output changed.
  */
-bool script_runner::set_output(const string&& new_output) {
+bool script_runner::set_output(string&& new_output) {
   std::lock_guard<std::mutex> guard(m_output_lock);
 
   if (m_output != new_output) {
@@ -98,16 +102,16 @@ script_runner::interval script_runner::run() {
     throw modules::module_error("Failed to execute command, stopping module...");
   }
 
-  int status = cmd->get_exit_status();
+  m_exit_status = cmd->get_exit_status();
   int fd = cmd->get_stdout(PIPE_READ);
   assert(fd != -1);
   bool changed = io_util::poll_read(fd) && set_output(cmd->readline());
 
-  if (!changed && status != 0) {
+  if (!changed && m_exit_status != 0) {
     clear_output();
   }
 
-  if (status == 0) {
+  if (m_exit_status == 0) {
     return m_interval;
   } else {
     return std::max(m_interval, interval{1s});
@@ -142,7 +146,7 @@ script_runner::interval script_runner::run_tail() {
     return 0s;
   }
 
-  bool exit_status = cmd->wait();
+  auto exit_status = cmd->wait();
 
   if (exit_status == 0) {
     return m_interval;
diff --git a/src/modules/script.cpp b/src/modules/script.cpp
index 71fe714a..af9466c9 100644
--- a/src/modules/script.cpp
+++ b/src/modules/script.cpp
@@ -25,7 +25,12 @@ namespace modules {
     // Setup formatting
     m_formatter->add(DEFAULT_FORMAT, TAG_LABEL, {TAG_LABEL});
     if (m_formatter->has(TAG_LABEL)) {
-      m_label = load_optional_label(m_conf, name(), "label", "%output%");
+      m_label = load_optional_label(m_conf, name(), TAG_LABEL, "%output%");
+    }
+
+    m_formatter->add_optional(FORMAT_FAIL, {TAG_LABEL_FAIL});
+    if (m_formatter->has(TAG_LABEL_FAIL)) {
+      m_label_fail = load_optional_label(m_conf, name(), TAG_LABEL_FAIL, "%output%");
     }
   }
 
@@ -70,6 +75,13 @@ namespace modules {
   /**
    * Generate module output
    */
+  string script_module::get_format() const {
+    if (m_runner.get_exit_status() != 0 && m_conf.has(name(), FORMAT_FAIL)) {
+      return FORMAT_FAIL;
+    }
+    return DEFAULT_FORMAT;
+  }
+
   string script_module::get_output() {
     auto script_output = m_runner.get_output();
     if (script_output.empty()) {
@@ -81,6 +93,11 @@ namespace modules {
       m_label->replace_token("%output%", script_output);
     }
 
+    if (m_label_fail) {
+      m_label_fail->reset_tokens();
+      m_label_fail->replace_token("%output%", script_output);
+    }
+
     string cnt{to_string(m_runner.get_counter())};
     string output{module::get_output()};
 
@@ -114,6 +131,8 @@ namespace modules {
   bool script_module::build(builder* builder, const string& tag) const {
     if (tag == TAG_LABEL) {
       builder->node(m_label);
+    } else if (tag == TAG_LABEL_FAIL) {
+      builder->node(m_label_fail);
     } else {
       return false;
     }