From 2ed95e235b11208331af9d03d1da69fdf44b47f9 Mon Sep 17 00:00:00 2001
From: Andreas Tsouchlos <andreas.tsouchlos@yardeye.com>
Date: Wed, 27 Apr 2022 14:13:49 +0200
Subject: [PATCH] Added Displayer::save_frames()

---
 display/dislpay_2d.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/display/dislpay_2d.py b/display/dislpay_2d.py
index cad35cf..fd5b209 100644
--- a/display/dislpay_2d.py
+++ b/display/dislpay_2d.py
@@ -44,7 +44,7 @@ class Displayer:
         self._objects.append(self._circle_map[obj_name])
 
 
-    # Functions to be used in the animation callback
+    # Functions to be used in the callback
 
 
     def move_object(self, obj_name, position):
@@ -63,7 +63,7 @@ class Displayer:
         self._circle_map[obj_name].set_alpha(1)
 
 
-    # Start the animation
+    # Display functions
 
 
     def animate(self, n_steps, anim_callback):
@@ -78,4 +78,16 @@ class Displayer:
                             interval=self._frame_interval,
                             init_func=self.__init_animation,
                             blit=False)
-        plt.show()
\ No newline at end of file
+        plt.show()
+    
+
+    def save_frames(self, frames, callback, prefix="fig-", ext="png"):
+        n_steps = max(frames) + 1
+
+        self.__init_animation()
+
+        for t in range(1, n_steps):
+            callback(self, t)
+
+            if t in frames:
+                self._fig.savefig(prefix + str(t) + "." + ext)