From 5b5a83b3162ee938c00fdba8f58753b8a594f1e3 Mon Sep 17 00:00:00 2001
From: Grant Sanderson <grant@3blue1brown.com>
Date: Sun, 12 Aug 2018 12:17:32 -0700
Subject: [PATCH] Added Mobject.update capabilties, which should be nicer and
 lighter weight than using ContinualUpdateFromFunc all the time.  Also changed
 the behavior of any Animation having its mobject automatically thrust up to
 the front of the screen

---
 animation/animation.py               |  2 -
 mobject/mobject.py                   | 43 +++++++++++++++-
 mobject/types/point_cloud_mobject.py |  2 +-
 scene/scene.py                       | 73 +++++++++++++++++++---------
 utils/simple_functions.py            |  5 ++
 5 files changed, 97 insertions(+), 28 deletions(-)

diff --git a/animation/animation.py b/animation/animation.py
index 410c6cc2..5797cacb 100644
--- a/animation/animation.py
+++ b/animation/animation.py
@@ -123,6 +123,4 @@ class Animation(object):
         if surrounding_scene is not None:
             if self.is_remover():
                 surrounding_scene.remove(self.mobject)
-            else:
-                surrounding_scene.add(self.mobject)
         return self
diff --git a/mobject/mobject.py b/mobject/mobject.py
index 23b3f75b..75b36fd8 100644
--- a/mobject/mobject.py
+++ b/mobject/mobject.py
@@ -1,6 +1,4 @@
 
-
-
 import copy
 import itertools as it
 import numpy as np
@@ -22,6 +20,7 @@ from utils.paths import straight_path
 from utils.space_ops import angle_of_vector
 from utils.space_ops import complex_to_R3
 from utils.space_ops import rotation_matrix
+from utils.simple_functions import get_num_args
 from functools import reduce
 
 
@@ -46,6 +45,7 @@ class Mobject(Container):
         self.color = Color(self.color)
         if self.name is None:
             self.name = self.__class__.__name__
+        self.updaters = []
         self.init_points()
         self.generate_points()
         self.init_colors()
@@ -98,6 +98,8 @@ class Mobject(Container):
             setattr(self, attr, func(getattr(self, attr)))
         return self
 
+    # Displaying
+
     def get_image(self, camera=None):
         if camera is None:
             from camera.camera import Camera
@@ -142,6 +144,42 @@ class Mobject(Container):
             self.target = self.copy()
         return self.target
 
+    # Updating
+
+    def update(self, dt):
+        for updater in self.updaters:
+            num_args = get_num_args(updater)
+            if num_args == 1:
+                updater(self)
+            elif num_args == 2:
+                updater(self, dt)
+            else:
+                raise Exception(
+                    "Mobject updater expected 1 or 2 "
+                    "arguments, %d given" % num_args
+                )
+
+    def get_time_based_updaters(self):
+        return [
+            udpater
+            for updater in self.updaters
+            if get_num_args(updater) == 2
+        ]
+
+    def get_updaters(self):
+        return self.updaters
+
+    def add_updater(self, update_function):
+        self.updaters.append(update_function)
+
+    def remove_updater(self, update_function):
+        while update_function in self.updaters:
+            self.updaters.remove(update_function)
+        return self
+
+    def clear_updaters(self):
+        self.updaters = []
+
     # Transforming operations
 
     def apply_to_family(self, func):
@@ -604,6 +642,7 @@ class Mobject(Container):
 
     def get_color(self):
         return self.color
+
     ##
 
     def save_state(self, use_deepcopy=False):
diff --git a/mobject/types/point_cloud_mobject.py b/mobject/types/point_cloud_mobject.py
index dbefff92..6677501c 100644
--- a/mobject/types/point_cloud_mobject.py
+++ b/mobject/types/point_cloud_mobject.py
@@ -13,7 +13,7 @@ from utils.iterables import stretch_array_to_length
 
 class PMobject(Mobject):
     CONFIG = {
-        "stroke_width": DEFAULT_POINT_THICKNESS,
+        "stroke_width": DEFAULT_STROKE_WIDTH,
     }
 
     def init_points(self):
diff --git a/scene/scene.py b/scene/scene.py
index 2cc50c35..9fee6196 100644
--- a/scene/scene.py
+++ b/scene/scene.py
@@ -172,9 +172,9 @@ class Scene(Container):
         self.clear()
     ###
 
-    def continual_update(self, dt=None):
-        if dt is None:
-            dt = self.frame_duration
+    def continual_update(self, dt):
+        for mobject in self.get_mobjects():
+            mobject.update(dt)
         for continual_animation in self.continual_animations:
             continual_animation.update(dt)
 
@@ -188,7 +188,17 @@ class Scene(Container):
         self.continual_animations = [ca for ca in self.continual_animations if ca in continual_animations]
 
     def should_continually_update(self):
-        return len(self.continual_animations) > 0 or self.always_continually_update
+        if self.always_continually_update:
+            return True
+        if len(self.continual_animations) > 0:
+            return True
+        any_time_based_update = any([
+            len(m.get_time_based_updaters()) > 0
+            for m in self.get_mobjects()
+        ])
+        if any_time_based_update:
+            return True
+        return False
 
     ###
 
@@ -230,6 +240,7 @@ class Scene(Container):
         mobjects, continual_animations = self.separate_mobjects_and_continual_animations(
             mobjects_or_continual_animations
         )
+        mobjects += self.foreground_mobjects
         self.restructure_mobjects(to_remove=mobjects)
         self.mobjects += mobjects
         self.continual_animations += continual_animations
@@ -330,15 +341,24 @@ class Scene(Container):
         return [m.copy() for m in self.mobjects]
 
     def get_moving_mobjects(self, *animations):
-        moving_mobjects = list(it.chain(
-            [
-                anim.mobject for anim in animations
-                if anim.mobject not in self.foreground_mobjects
-            ],
-            [ca.mobject for ca in self.continual_animations],
-            self.foreground_mobjects,
-        ))
-        return moving_mobjects
+        # Go through mobjects from start to end, and
+        # as soon as there's one that needs updating of
+        # some kind per frame, return the list from that
+        # point forward.
+        animation_mobjects = [anim.mobject for anim in animations]
+        ca_mobjects = [ca.mobject for ca in self.continual_animations]
+        mobjects = self.get_mobjects()
+        for i, mob in enumerate(mobjects):
+            update_possibilities = [
+                mob in animation_mobjects,
+                mob in ca_mobjects,
+                len(mob.get_updaters()) > 0,
+                mob in self.foreground_mobjects
+            ]
+            for possibility in update_possibilities:
+                if possibility:
+                    return mobjects[i:]
+        return []
 
     def get_time_progression(self, run_time):
         if self.skip_animations:
@@ -439,25 +459,30 @@ class Scene(Container):
             # This is where kwargs to play like run_time and rate_func
             # get applied to all animations
             animation.update_config(**kwargs)
+            # Anything animated that's not already in the
+            # scene gets added to the scene
+            if animation.mobject not in self.mobjects:
+                self.add(animation.mobject)
         moving_mobjects = self.get_moving_mobjects(*animations)
 
         # Paint all non-moving objects onto the screen, so they don't
         # have to be rendered every frame
         self.update_frame(excluded_mobjects=moving_mobjects)
         static_image = self.get_frame()
+        total_run_time = 0
         for t in self.get_animation_time_progression(animations):
+            self.continual_update(dt=t - total_run_time)
             for animation in animations:
                 animation.update(t / animation.run_time)
-            self.continual_update()
             self.update_frame(moving_mobjects, static_image)
             self.add_frames(self.get_frame())
-        self.add(*moving_mobjects)
-        self.mobjects_from_last_animation = moving_mobjects
+            total_run_time = t
+        self.mobjects_from_last_animation = [
+            anim.mobject for anim in animations
+        ]
         self.clean_up_animations(*animations)
         if self.skip_animations:
-            # Todo, not great that this uses a variable from
-            # a previous loop...
-            self.continual_update(t)
+            self.continual_update(total_run_time)
         else:
             self.continual_update(0)
         self.num_plays += 1
@@ -466,7 +491,6 @@ class Scene(Container):
     def clean_up_animations(self, *animations):
         for animation in animations:
             animation.clean_up(self)
-        self.add(*self.foreground_mobjects)
         return self
 
     def get_mobjects_from_last_animation(self):
@@ -476,17 +500,20 @@ class Scene(Container):
 
     def wait(self, duration=DEFAULT_WAIT_TIME):
         if self.should_continually_update():
+            total_time = 0
             for t in self.get_time_progression(duration):
-                self.continual_update()
+                self.continual_update(dt=t - total_time)
                 self.update_frame()
                 self.add_frames(self.get_frame())
+                total_time = t
         elif self.skip_animations:
             # Do nothing
             return self
         else:
             self.update_frame()
-            self.add_frames(*[self.get_frame()] *
-                            int(duration / self.frame_duration))
+            n_frames = int(duration / self.frame_duration)
+            frame = self.get_frame()
+            self.add_frames(*[frame] * n_frames)
         return self
 
     def wait_to(self, time, assert_positive=True):
diff --git a/utils/simple_functions.py b/utils/simple_functions.py
index bf4540d6..48f88e4f 100644
--- a/utils/simple_functions.py
+++ b/utils/simple_functions.py
@@ -1,5 +1,6 @@
 import numpy as np
 import operator as op
+import inspect
 from functools import reduce
 
 
@@ -16,6 +17,10 @@ def choose(n, r):
     numer = reduce(op.mul, range(n, n - r, -1), 1)
     return numer // denom
 
+
+def get_num_args(function):
+    return len(inspect.signature(function).parameters)
+
 # Just to have a less heavyweight name for this extremely common operation
 #
 # We may wish to have more fine-grained control over division by zero behavior