From c7fa8c59adf7d18e4591e398c5418ab58a6dde92 Mon Sep 17 00:00:00 2001 From: Grant Sanderson Date: Wed, 24 Jan 2018 11:25:55 -0800 Subject: [PATCH] Made slicing into mobjects return a Group or VGroup instead of a list --- mobject/mobject.py | 19 +++++++++++++------ mobject/vectorized_mobject.py | 3 +++ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/mobject/mobject.py b/mobject/mobject.py index c8bfc5d3..f079b330 100644 --- a/mobject/mobject.py +++ b/mobject/mobject.py @@ -15,11 +15,11 @@ class Mobject(object): Mathematical Object """ CONFIG = { - "color" : WHITE, + "color" : WHITE, "stroke_width" : DEFAULT_POINT_THICKNESS, - "name" : None, - "dim" : 3, - "target" : None, + "name" : None, + "dim" : 3, + "target" : None, } def __init__(self, *submobjects, **kwargs): digest_config(self, kwargs) @@ -663,8 +663,12 @@ class Mobject(object): ## Family matters - def __getitem__(self, index): - return self.split()[index] + def __getitem__(self, value): + self_list = self.split() + if isinstance(value, slice): + GroupClass = self.get_group_class() + return GroupClass(*self_list.__getitem__(value)) + return self_list.__getitem__(value) def __iter__(self): return iter(self.split()) @@ -672,6 +676,9 @@ class Mobject(object): def __len__(self): return len(self.split()) + def get_group_class(self): + return Group + def split(self): result = [self] if len(self.points) > 0 else [] return result + self.submobjects diff --git a/mobject/vectorized_mobject.py b/mobject/vectorized_mobject.py index 264181fd..4b1a37e6 100644 --- a/mobject/vectorized_mobject.py +++ b/mobject/vectorized_mobject.py @@ -19,6 +19,9 @@ class VMobject(Mobject): "make_smooth_after_applying_functions" : False, } + def get_group_class(self): + return VGroup + ## Colors def init_colors(self): self.set_style_data(