From a0bd95a790af1169eea0399782bb6b590a7d7302 Mon Sep 17 00:00:00 2001
From: Long Vo <hoanglonguk@gmail.com>
Date: Tue, 19 Nov 2024 14:17:58 +0700
Subject: [PATCH] Added mqtt config to publish nan as none

---
 esphome/components/mqtt/__init__.py     | 4 ++++
 esphome/components/mqtt/mqtt_client.cpp | 2 ++
 esphome/components/mqtt/mqtt_client.h   | 6 ++++++
 esphome/components/mqtt/mqtt_sensor.cpp | 2 +-
 esphome/const.py                        | 1 +
 5 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/esphome/components/mqtt/__init__.py b/esphome/components/mqtt/__init__.py
index 86d163e61d..12024820c4 100644
--- a/esphome/components/mqtt/__init__.py
+++ b/esphome/components/mqtt/__init__.py
@@ -49,6 +49,7 @@ from esphome.const import (
     CONF_USE_ABBREVIATIONS,
     CONF_USERNAME,
     CONF_WILL_MESSAGE,
+    CONF_PUBLISH_NAN_AS_NONE,
     PLATFORM_BK72XX,
     PLATFORM_ESP32,
     PLATFORM_ESP8266,
@@ -296,6 +297,7 @@ CONFIG_SCHEMA = cv.All(
                     cv.Optional(CONF_QOS, default=0): cv.mqtt_qos,
                 }
             ),
+            cv.Optional(CONF_PUBLISH_NAN_AS_NONE): cv.boolean,
         }
     ),
     validate_config,
@@ -449,6 +451,8 @@ async def to_code(config):
         trigger = cg.new_Pvariable(conf[CONF_TRIGGER_ID], var)
         await automation.build_automation(trigger, [], conf)
 
+    if CONF_PUBLISH_NAN_AS_NONE in config:
+        cg.add(var.set_publish_nan_as_none(config[CONF_PUBLISH_NAN_AS_NONE]))
 
 MQTT_PUBLISH_ACTION_SCHEMA = cv.Schema(
     {
diff --git a/esphome/components/mqtt/mqtt_client.cpp b/esphome/components/mqtt/mqtt_client.cpp
index 106192c0e3..9139dfdc00 100644
--- a/esphome/components/mqtt/mqtt_client.cpp
+++ b/esphome/components/mqtt/mqtt_client.cpp
@@ -608,6 +608,8 @@ void MQTTClientComponent::set_log_message_template(MQTTMessage &&message) { this
 const MQTTDiscoveryInfo &MQTTClientComponent::get_discovery_info() const { return this->discovery_info_; }
 void MQTTClientComponent::set_topic_prefix(const std::string &topic_prefix) { this->topic_prefix_ = topic_prefix; }
 const std::string &MQTTClientComponent::get_topic_prefix() const { return this->topic_prefix_; }
+void MQTTClientComponent::set_publish_nan_as_none(bool publish_nan_as_none) { this->publish_nan_as_none_ = publish_nan_as_none; }
+bool MQTTClientComponent::is_publish_nan_as_none() const { return this->publish_nan_as_none_; }
 void MQTTClientComponent::disable_birth_message() {
   this->birth_message_.topic = "";
   this->recalculate_availability_();
diff --git a/esphome/components/mqtt/mqtt_client.h b/esphome/components/mqtt/mqtt_client.h
index 7ae3a6c5e8..80a4ef51a3 100644
--- a/esphome/components/mqtt/mqtt_client.h
+++ b/esphome/components/mqtt/mqtt_client.h
@@ -263,6 +263,10 @@ class MQTTClientComponent : public Component {
   void set_on_connect(mqtt_on_connect_callback_t &&callback);
   void set_on_disconnect(mqtt_on_disconnect_callback_t &&callback);
 
+  // Publish None state instead of NaN for Home Assistant
+  void set_publish_nan_as_none(bool publish_nan_as_none_);
+  bool is_publish_nan_as_none() const;
+
  protected:
   void send_device_info_();
 
@@ -328,6 +332,8 @@ class MQTTClientComponent : public Component {
   uint32_t connect_begin_;
   uint32_t last_connected_{0};
   optional<MQTTClientDisconnectReason> disconnect_reason_{};
+
+  bool publish_nan_as_none_{false};
 };
 
 extern MQTTClientComponent *global_mqtt_client;  // NOLINT(cppcoreguidelines-avoid-non-const-global-variables)
diff --git a/esphome/components/mqtt/mqtt_sensor.cpp b/esphome/components/mqtt/mqtt_sensor.cpp
index 50c41beaa1..2cbc291ccf 100644
--- a/esphome/components/mqtt/mqtt_sensor.cpp
+++ b/esphome/components/mqtt/mqtt_sensor.cpp
@@ -69,7 +69,7 @@ bool MQTTSensorComponent::send_initial_state() {
   }
 }
 bool MQTTSensorComponent::publish_state(float value) {
-  if (std::isnan(value))
+  if (mqtt::global_mqtt_client->is_publish_nan_as_none() && std::isnan(value))
     return this->publish(this->get_state_topic_(), "None");
   int8_t accuracy = this->sensor_->get_accuracy_decimals();
   return this->publish(this->get_state_topic_(), value_accuracy_to_string(value, accuracy));
diff --git a/esphome/const.py b/esphome/const.py
index 6a643e1e30..1546346898 100644
--- a/esphome/const.py
+++ b/esphome/const.py
@@ -955,6 +955,7 @@ CONF_X_GRID = "x_grid"
 CONF_Y_GRID = "y_grid"
 CONF_YEAR = "year"
 CONF_ZERO = "zero"
+CONF_PUBLISH_NAN_AS_NONE = "publish_nan_as_none"
 
 TYPE_GIT = "git"
 TYPE_LOCAL = "local"