From 9cd173ef83b4a0acfc42a5406725d81d064fcb6a Mon Sep 17 00:00:00 2001
From: guillempages <guillempages@users.noreply.github.com>
Date: Thu, 25 May 2023 23:49:52 +0200
Subject: [PATCH] Allow partially looping animations (#4693)

Add the possibility of specifying a "loop" in an animation; where the
requested frames (start - end) will be repeateadly shown for "count" times.
---
 esphome/components/animation/__init__.py      | 25 ++++++++++++++++++-
 esphome/components/display/display_buffer.cpp | 23 +++++++++++++++--
 esphome/components/display/display_buffer.h   | 10 ++++++--
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/esphome/components/animation/__init__.py b/esphome/components/animation/__init__.py
index 1b804bd527..f51d115d9e 100644
--- a/esphome/components/animation/__init__.py
+++ b/esphome/components/animation/__init__.py
@@ -6,7 +6,14 @@ import esphome.components.image as espImage
 from esphome.components.image import CONF_USE_TRANSPARENCY
 import esphome.config_validation as cv
 import esphome.codegen as cg
-from esphome.const import CONF_FILE, CONF_ID, CONF_RAW_DATA_ID, CONF_RESIZE, CONF_TYPE
+from esphome.const import (
+    CONF_FILE,
+    CONF_ID,
+    CONF_RAW_DATA_ID,
+    CONF_REPEAT,
+    CONF_RESIZE,
+    CONF_TYPE,
+)
 from esphome.core import CORE, HexInt
 
 _LOGGER = logging.getLogger(__name__)
@@ -14,6 +21,10 @@ _LOGGER = logging.getLogger(__name__)
 DEPENDENCIES = ["display"]
 MULTI_CONF = True
 
+CONF_LOOP = "loop"
+CONF_START_FRAME = "start_frame"
+CONF_END_FRAME = "end_frame"
+
 Animation_ = display.display_ns.class_("Animation", espImage.Image_)
 
 
@@ -48,6 +59,13 @@ ANIMATION_SCHEMA = cv.Schema(
             # Not setting default here on purpose; the default depends on the image type,
             # and thus will be set in the "validate_cross_dependencies" validator.
             cv.Optional(CONF_USE_TRANSPARENCY): cv.boolean,
+            cv.Optional(CONF_LOOP): cv.All(
+                {
+                    cv.Optional(CONF_START_FRAME, default=0): cv.positive_int,
+                    cv.Optional(CONF_END_FRAME): cv.positive_int,
+                    cv.Optional(CONF_REPEAT): cv.positive_int,
+                }
+            ),
             cv.GenerateID(CONF_RAW_DATA_ID): cv.declare_id(cg.uint8),
         },
         validate_cross_dependencies,
@@ -227,3 +245,8 @@ async def to_code(config):
         espImage.IMAGE_TYPE[config[CONF_TYPE]],
     )
     cg.add(var.set_transparency(transparent))
+    if CONF_LOOP in config:
+        start = config[CONF_LOOP][CONF_START_FRAME]
+        end = config[CONF_LOOP].get(CONF_END_FRAME, frames)
+        count = config[CONF_LOOP].get(CONF_REPEAT, -1)
+        cg.add(var.set_loop(start, end, count))
diff --git a/esphome/components/display/display_buffer.cpp b/esphome/components/display/display_buffer.cpp
index 35e55bc1ba..0d76fa09ec 100644
--- a/esphome/components/display/display_buffer.cpp
+++ b/esphome/components/display/display_buffer.cpp
@@ -773,12 +773,31 @@ Color Animation::get_grayscale_pixel(int x, int y) const {
   return Color(gray, gray, gray, alpha);
 }
 Animation::Animation(const uint8_t *data_start, int width, int height, uint32_t animation_frame_count, ImageType type)
-    : Image(data_start, width, height, type), current_frame_(0), animation_frame_count_(animation_frame_count) {}
-int Animation::get_animation_frame_count() const { return this->animation_frame_count_; }
+    : Image(data_start, width, height, type),
+      current_frame_(0),
+      animation_frame_count_(animation_frame_count),
+      loop_start_frame_(0),
+      loop_end_frame_(animation_frame_count_),
+      loop_count_(0),
+      loop_current_iteration_(1) {}
+void Animation::set_loop(uint32_t start_frame, uint32_t end_frame, int count) {
+  loop_start_frame_ = std::min(start_frame, animation_frame_count_);
+  loop_end_frame_ = std::min(end_frame, animation_frame_count_);
+  loop_count_ = count;
+  loop_current_iteration_ = 1;
+}
+
+uint32_t Animation::get_animation_frame_count() const { return this->animation_frame_count_; }
 int Animation::get_current_frame() const { return this->current_frame_; }
 void Animation::next_frame() {
   this->current_frame_++;
+  if (loop_count_ && this->current_frame_ == loop_end_frame_ &&
+      (this->loop_current_iteration_ < loop_count_ || loop_count_ < 0)) {
+    this->current_frame_ = loop_start_frame_;
+    this->loop_current_iteration_++;
+  }
   if (this->current_frame_ >= animation_frame_count_) {
+    this->loop_current_iteration_ = 1;
     this->current_frame_ = 0;
   }
 }
diff --git a/esphome/components/display/display_buffer.h b/esphome/components/display/display_buffer.h
index a8ec0e588f..2474d6f5a0 100644
--- a/esphome/components/display/display_buffer.h
+++ b/esphome/components/display/display_buffer.h
@@ -569,7 +569,7 @@ class Animation : public Image {
   Color get_rgb565_pixel(int x, int y) const override;
   Color get_grayscale_pixel(int x, int y) const override;
 
-  int get_animation_frame_count() const;
+  uint32_t get_animation_frame_count() const;
   int get_current_frame() const override;
   void next_frame();
   void prev_frame();
@@ -580,9 +580,15 @@ class Animation : public Image {
    */
   void set_frame(int frame);
 
+  void set_loop(uint32_t start_frame, uint32_t end_frame, int count);
+
  protected:
   int current_frame_;
-  int animation_frame_count_;
+  uint32_t animation_frame_count_;
+  uint32_t loop_start_frame_;
+  uint32_t loop_end_frame_;
+  int loop_count_;
+  int loop_current_iteration_;
 };
 
 template<typename... Ts> class DisplayPageShowAction : public Action<Ts...> {