From 8b1f0a8749d91eeda4b674ed156cbc7f8e1e48a8 Mon Sep 17 00:00:00 2001 From: Grant Sanderson Date: Tue, 22 Mar 2022 10:35:49 -0700 Subject: [PATCH] Refactor Mobject.set_rgba_array_by_color --- manimlib/mobject/mobject.py | 34 +++++++++++++--------------------- 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/manimlib/mobject/mobject.py b/manimlib/mobject/mobject.py index cf1889bf..d8f7bec8 100644 --- a/manimlib/mobject/mobject.py +++ b/manimlib/mobject/mobject.py @@ -935,32 +935,24 @@ class Mobject(object): return self def set_rgba_array_by_color(self, color=None, opacity=None, name="rgbas", recurse=True): + max_len = 0 if color is not None: rgbs = np.array([color_to_rgb(c) for c in listify(color)]) + max_len = len(rgbs) if opacity is not None: - opacities = listify(opacity) + opacities = np.array(listify(opacity)) + max_len = max(max_len, len(opacities)) - # Color only - if color is not None and opacity is None: - for mob in self.get_family(recurse): - mob.data[name] = resize_array(mob.data[name], len(rgbs)) - mob.data[name][:, :3] = rgbs - - # Opacity only - if color is None and opacity is not None: - for mob in self.get_family(recurse): - mob.data[name] = resize_array(mob.data[name], len(opacities)) - mob.data[name][:, 3] = opacities - - # Color and opacity - if color is not None and opacity is not None: - rgbas = np.array([ - [*rgb, o] - for rgb, o in zip(*make_even(rgbs, opacities)) - ]) - for mob in self.get_family(recurse): - mob.data[name] = rgbas.copy() + for mob in self.get_family(recurse): + if max_len > len(mob.data[name]): + mob.data[name] = resize_array(mob.data[name], max_len) + size = len(mob.data[name]) + if color is not None: + mob.data[name][:, :3] = resize_array(rgbs, size) + if opacity is not None: + mob.data[name][:, 3] = resize_array(opacities, size) return self + def set_color(self, color, opacity=None, recurse=True): self.set_rgba_array_by_color(color, opacity, recurse=False)