Extended DecimalNumber to handles complex numbers and background rectangles

topics/numerals.py

@@ -3,6 +3,7 @@ from mobject.vectorized_mobject import VMobject, VGroup, VectorizedPoint
 from mobject.tex_mobject import TexMobject
 from animation import Animation
 from animation.continual_animation import ContinualAnimation
+from topics.geometry import BackgroundRectangle
 from scene import Scene
 from helpers import *
 
@@ -12,18 +13,31 @@ class DecimalNumber(VMobject):
         "digit_to_digit_buff" : 0.05,
         "show_ellipsis" : False,
         "unit" : None,
+        "include_background_rectangle" : False,
     }
     def __init__(self, number, **kwargs):
-        digest_config(self, kwargs, locals())
-        num_string = '%.*f'%(self.num_decimal_points, number)
-        negative_zero_string = "-%.*f"%(self.num_decimal_points, 0.)
-        if num_string == negative_zero_string:
-            num_string = num_string[1:]
-        VMobject.__init__(self, *[
-            TexMobject(char)
-            for char in num_string
-        ], **kwargs)
+        VMobject.__init__(self, **kwargs)
+        self.number = number
+        ndp = self.num_decimal_points
 
+        #Build number string
+        if isinstance(number, complex):
+            num_string = '%.*f%s%.*fi'%(
+                ndp, number.real, 
+                "-" if number.imag < 0 else "+",
+                ndp, abs(number.imag)
+            )
+        else:
+            num_string = '%.*f'%(ndp, number)
+            negative_zero_string = "-%.*f"%(ndp, 0.)
+            if num_string == negative_zero_string:
+                num_string = num_string[1:]
+        self.add(*[
+            TexMobject(char, **kwargs)
+            for char in num_string
+        ])
+
+        #Add non-numerical bits
         if self.show_ellipsis:
             self.add(TexMobject("\\dots"))
 
@@ -35,28 +49,27 @@ class DecimalNumber(VMobject):
             aligned_edge = DOWN
         )
 
-        if num_string.startswith("-"):
-            minus = self.submobjects[0]
-            minus.next_to(
-                self.submobjects[1], LEFT,
-                buff = self.digit_to_digit_buff
-            )
+        #Handle alignment of parts that should be aligned
+        #to the bottom
+        for i, c in enumerate(num_string):
+            if c == "-" and len(num_string) > i+1:
+                self[i].align_to(self[i+1], alignment_vect = UP)
         if self.unit == "\\circ":
             self[-1].align_to(self, UP)
+        #
+        if self.include_background_rectangle:
+            #TODO, is this the best way to handle
+            #background rectangles?
+            self.background_rectangle = BackgroundRectangle(self)
+            self.submobjects = [
+                self.background_rectangle,
+                VGroup(*self.submobjects)
+            ]
 
-class Integer(VGroup):
+class Integer(DecimalNumber):
     CONFIG = {
-        "digit_buff" : 0.8*SMALL_BUFF
+        "num_decimal_points" : 0,
     }
-    def __init__(self, integer, **kwargs):
-        self.number = integer
-        num_str = str(integer)
-        VGroup.__init__(self, *map(TexMobject, num_str), **kwargs)
-        self.arrange_submobjects(
-            RIGHT, buff = self.digit_buff, aligned_edge = DOWN
-        )
-        if num_str[0] == "-":
-            self[0].next_to(self[1], LEFT, buff = SMALL_BUFF)
 
 class ChangingDecimal(Animation):
     CONFIG = {