Adding new blocks to existing animations

animatplot/animation.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from animatplot import Timeline
+from animatplot.blocks.base import Block
 
 
 class Animation:
@@ -26,6 +27,11 @@ class Animation:
         a matplotlib animation returned from FuncAnimation
     """
     def __init__(self, blocks, timeline=None, fig=None):
+        self.fig = plt.gcf() if fig is None else fig
+
+        self.animation = self._animate(blocks, timeline)
+
+    def _animate(self, blocks, timeline):
         if timeline is None:
             self.timeline = Timeline(range(len(blocks[0])))
         elif not isinstance(timeline, Timeline):
@@ -36,14 +42,14 @@ def __init__(self, blocks, timeline=None, fig=None):
         _len_time = len(self.timeline)
         for block in blocks:
             if len(block) != _len_time:
-                raise ValueError("All blocks must animate for the same amount of time")
+                raise ValueError("All blocks must animate for the same amount "
+                                 "of time")
 
         self.blocks = blocks
-        self.fig = plt.gcf() if fig is None else fig
         self._has_slider = False
         self._pause = False
 
-        def animate(i):
+        def update_all(i):
             updates = []
             for block in self.blocks:
                 updates.append(block._update(self.timeline.index))
@@ -52,11 +58,8 @@ def animate(i):
             self.timeline._update()
             return updates
 
-        self.animation = FuncAnimation(
-            self.fig, animate,
-            frames=self.timeline._len,
-            interval=1000/self.timeline.fps
-        )
+        return FuncAnimation(self.fig, update_all, frames=self.timeline._len,
+                             interval=1000 / self.timeline.fps)
 
     def toggle(self, ax=None):
         """Creates a play/pause button to start/stop the animation
@@ -176,7 +179,8 @@ def save_gif(self, filename):
             the name of the file to be created without the file extension
         """
         self.timeline.index -= 1  # required for proper starting point for save
-        self.animation.save(filename+'.gif', writer=PillowWriter(fps=self.timeline.fps))
+        self.animation.save(filename+'.gif',
+                            writer=PillowWriter(fps=self.timeline.fps))
 
     def save(self, *args, **kwargs):
         """Saves an animation
@@ -185,3 +189,39 @@ def save(self, *args, **kwargs):
         """
         self.timeline.index -= 1  # required for proper starting point for save
         self.animation.save(*args, **kwargs)
+
+    def add(self, new):
+        """
+        Updates the animation object by adding additional blocks.
+
+        The new blocks can be passed as a list, or as part of a second animaion.
+        If passed as part of a new animation, the timeline of this new
+        animation object will replace the old one.
+
+        Parameters
+        ----------
+        new : amp.animation.Animation, or list of amp.block.Block objects
+            Either blocks to add to animation instance, or another animation
+            instance whose blocks should be combined with this animation.
+        """
+
+        if isinstance(new, Animation):
+            new_blocks = new.blocks
+            new_timeline = new.timeline
+
+        else:
+            if not isinstance(new, list):
+                new_blocks = [new]
+            else:
+                new_blocks = new
+            new_timeline = self.timeline
+
+        for i, block in enumerate(new_blocks):
+            if not isinstance(block, Block):
+                raise TypeError(f"Block number {i} passed is of type "
+                                f"{type(block)}, not of type "
+                                f"animatplot.blocks.Block (or a subclass)")
+
+            self.blocks.append(block)
+
+        self.animation = self._animate(self.blocks, new_timeline)

tests/test_animation.py

@@ -1,14 +1,21 @@
-from matplotlib.testing import setup
-setup()
 import os
+
 import pytest
-import numpy as np
+
+import matplotlib.testing
 import matplotlib.pyplot as plt
 from matplotlib.animation import PillowWriter
+
+import numpy as np
+import numpy.testing as npt
+
 import animatplot as amp
 from tests.tools import animation_compare
 
 
+matplotlib.testing.setup()
+
+
 @pytest.mark.xfail
 @animation_compare(baseline_images='Animation/controls', nframes=5, tol=.5)
 def test_controls():
@@ -41,3 +48,90 @@ def test_save():
     anim.save_gif(base+'save')
     plt.close('all')
     assert os.path.exists(base+'save.gif')
+
+
+@pytest.fixture()
+def line_block():
+    def make_line_block(t_length=5):
+        x = np.linspace(0, 1, 10)
+        t = np.linspace(0, 1, t_length)
+        x_grid, t_grid = np.meshgrid(x, t)
+        y_data = np.sin(2 * np.pi * (x_grid + t_grid))
+
+        return amp.blocks.Line(x, y_data)
+    return make_line_block
+
+
+@pytest.fixture()
+def line_anim():
+    def make_line_anim(t_length=5, timeline=False):
+        x = np.linspace(0, 1, 10)
+        t = np.linspace(0, 1, t_length)
+        x_grid, t_grid = np.meshgrid(x, t)
+        y_data = np.sin(2 * np.pi * (x_grid + t_grid))
+
+        block = amp.blocks.Line(x, y_data)
+
+        if timeline:
+            return amp.Animation([block], timeline=amp.Timeline(t))
+        else:
+            return amp.Animation([block])
+
+    return make_line_anim
+
+
+class TestAddBlocks:
+    def test_add_blocks(self, line_block):
+        anim = amp.Animation([line_block()])
+        anim.add(line_block())
+
+        assert isinstance(anim, amp.Animation)
+        assert len(anim.blocks) == 2
+        for actual in anim.blocks:
+            assert len(actual) == 5
+            npt.assert_equal(actual.line.get_xdata(),
+                             np.linspace(0, 1, 10))
+
+    def test_wrong_length_block(self, line_block):
+        anim = amp.Animation([line_block()])
+
+        with pytest.raises(ValueError):
+            anim.add(line_block(t_length=6))
+
+    def test_wrong_type(self, line_block):
+        anim = amp.Animation([line_block()])
+
+        with pytest.raises(TypeError):
+            anim.add('not a block')
+
+
+class TestAddAnimations:
+    def test_add_animations(self, line_anim):
+        anim = line_anim()
+        anim.add(line_anim())
+
+        assert isinstance(anim, amp.Animation)
+
+    def test_add_animation_with_timeline(self, line_anim):
+        anim = line_anim()
+        anim.add(line_anim(timeline=True))
+
+        assert isinstance(anim, amp.Animation)
+        assert len(anim.timeline) == 5
+
+    def test_add_animations_both_with_timelines(self, line_anim):
+        anim = line_anim(timeline=True)
+        anim2 = line_anim()
+        t = 10*np.arange(5)
+        anim2.timeline = amp.Timeline(t)
+
+        anim.add(anim2)
+
+        assert isinstance(anim, amp.Animation)
+        assert len(anim.timeline) == 5
+        npt.assert_equal(anim.timeline.t, t)
+
+    def test_add_animations_different_lengths(self, line_anim):
+        anim = line_anim()
+        with pytest.raises(ValueError):
+            anim.add(line_anim(t_length=6))