From 98f0d75180e7eedb67a41b481427e7beef0472e7 Mon Sep 17 00:00:00 2001
From: 0x0a11c0de <77504775+0x0a11c0de@users.noreply.github.com>
Date: Mon, 3 May 2021 19:11:24 -0700
Subject: [PATCH] Fix #1940: Implement speed_count in TuyaFan (#1654)

Co-authored-by: Frank Riley <fhriley@gmail.com>
---
 esphome/components/tuya/fan/__init__.py  | 15 ++++++++-------
 esphome/components/tuya/fan/tuya_fan.cpp | 12 +++++-------
 esphome/components/tuya/fan/tuya_fan.h   |  4 ++--
 3 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/esphome/components/tuya/fan/__init__.py b/esphome/components/tuya/fan/__init__.py
index 8615f3ae85..b5e3e579be 100644
--- a/esphome/components/tuya/fan/__init__.py
+++ b/esphome/components/tuya/fan/__init__.py
@@ -1,7 +1,7 @@
 from esphome.components import fan
 import esphome.config_validation as cv
 import esphome.codegen as cg
-from esphome.const import CONF_OUTPUT_ID, CONF_SWITCH_DATAPOINT
+from esphome.const import CONF_OUTPUT_ID, CONF_SPEED_COUNT, CONF_SWITCH_DATAPOINT
 from .. import tuya_ns, CONF_TUYA_ID, Tuya
 
 DEPENDENCIES = ["tuya"]
@@ -19,6 +19,7 @@ CONFIG_SCHEMA = cv.All(
             cv.Optional(CONF_OSCILLATION_DATAPOINT): cv.uint8_t,
             cv.Optional(CONF_SPEED_DATAPOINT): cv.uint8_t,
             cv.Optional(CONF_SWITCH_DATAPOINT): cv.uint8_t,
+            cv.Optional(CONF_SPEED_COUNT, default=3): cv.int_range(min=1, max=256),
         }
     ).extend(cv.COMPONENT_SCHEMA),
     cv.has_at_least_one_key(CONF_SPEED_DATAPOINT, CONF_SWITCH_DATAPOINT),
@@ -26,13 +27,13 @@ CONFIG_SCHEMA = cv.All(
 
 
 def to_code(config):
-    var = cg.new_Pvariable(config[CONF_OUTPUT_ID])
-    yield cg.register_component(var, config)
+    parent = yield cg.get_variable(config[CONF_TUYA_ID])
+    state = yield fan.create_fan_state(config)
 
-    paren = yield cg.get_variable(config[CONF_TUYA_ID])
-    fan_ = yield fan.create_fan_state(config)
-    cg.add(var.set_tuya_parent(paren))
-    cg.add(var.set_fan(fan_))
+    var = cg.new_Pvariable(
+        config[CONF_OUTPUT_ID], parent, state, config[CONF_SPEED_COUNT]
+    )
+    yield cg.register_component(var, config)
 
     if CONF_SPEED_DATAPOINT in config:
         cg.add(var.set_speed_id(config[CONF_SPEED_DATAPOINT]))
diff --git a/esphome/components/tuya/fan/tuya_fan.cpp b/esphome/components/tuya/fan/tuya_fan.cpp
index 718f292f5b..26e5de253d 100644
--- a/esphome/components/tuya/fan/tuya_fan.cpp
+++ b/esphome/components/tuya/fan/tuya_fan.cpp
@@ -8,18 +8,15 @@ namespace tuya {
 static const char *TAG = "tuya.fan";
 
 void TuyaFan::setup() {
-  auto traits = fan::FanTraits(this->oscillation_id_.has_value(), this->speed_id_.has_value(), false, 3);
+  auto traits =
+      fan::FanTraits(this->oscillation_id_.has_value(), this->speed_id_.has_value(), false, this->speed_count_);
   this->fan_->set_traits(traits);
 
   if (this->speed_id_.has_value()) {
     this->parent_->register_listener(*this->speed_id_, [this](TuyaDatapoint datapoint) {
       auto call = this->fan_->make_call();
-      if (datapoint.value_enum == 0x0)
-        call.set_speed(1);
-      else if (datapoint.value_enum == 0x1)
-        call.set_speed(2);
-      else if (datapoint.value_enum == 0x2)
-        call.set_speed(3);
+      if (datapoint.value_enum < this->speed_count_)
+        call.set_speed(datapoint.value_enum + 1);
       else
         ESP_LOGCONFIG(TAG, "Speed has invalid value %d", datapoint.value_enum);
       ESP_LOGD(TAG, "MCU reported speed of: %d", datapoint.value_enum);
@@ -47,6 +44,7 @@ void TuyaFan::setup() {
 
 void TuyaFan::dump_config() {
   ESP_LOGCONFIG(TAG, "Tuya Fan:");
+  ESP_LOGCONFIG(TAG, "  Speed count %d", this->speed_count_);
   if (this->speed_id_.has_value())
     ESP_LOGCONFIG(TAG, "  Speed has datapoint ID %u", *this->speed_id_);
   if (this->switch_id_.has_value())
diff --git a/esphome/components/tuya/fan/tuya_fan.h b/esphome/components/tuya/fan/tuya_fan.h
index d31d490e1a..ae31df087c 100644
--- a/esphome/components/tuya/fan/tuya_fan.h
+++ b/esphome/components/tuya/fan/tuya_fan.h
@@ -9,13 +9,12 @@ namespace tuya {
 
 class TuyaFan : public Component {
  public:
+  TuyaFan(Tuya *parent, fan::FanState *fan, int speed_count) : parent_(parent), fan_(fan), speed_count_(speed_count) {}
   void setup() override;
   void dump_config() override;
   void set_speed_id(uint8_t speed_id) { this->speed_id_ = speed_id; }
   void set_switch_id(uint8_t switch_id) { this->switch_id_ = switch_id; }
   void set_oscillation_id(uint8_t oscillation_id) { this->oscillation_id_ = oscillation_id; }
-  void set_fan(fan::FanState *fan) { this->fan_ = fan; }
-  void set_tuya_parent(Tuya *parent) { this->parent_ = parent; }
   void write_state();
 
  protected:
@@ -28,6 +27,7 @@ class TuyaFan : public Component {
   optional<uint8_t> switch_id_{};
   optional<uint8_t> oscillation_id_{};
   fan::FanState *fan_;
+  int speed_count_{};
 };
 
 }  // namespace tuya