Add surround method for Rectangles and SurroundingRectangles

manimlib/mobject/geometry.py

@@ -7,7 +7,7 @@ import numpy as np
 
 from manimlib.constants import DL, DOWN, DR, LEFT, ORIGIN, OUT, RIGHT, UL, UP, UR
 from manimlib.constants import GREY_A, RED, WHITE, BLACK
-from manimlib.constants import MED_SMALL_BUFF
+from manimlib.constants import MED_SMALL_BUFF, SMALL_BUFF
 from manimlib.constants import DEGREES, PI, TAU
 from manimlib.mobject.mobject import Mobject
 from manimlib.mobject.types.vectorized_mobject import DashedVMobject
@@ -1046,6 +1046,12 @@ class Rectangle(Polygon):
         self.set_width(width, stretch=True)
         self.set_height(height, stretch=True)
 
+    def surround(self, mobject, buff=SMALL_BUFF) -> Self:
+        target_shape = np.array(mobject.get_shape()) + 2 * buff
+        self.set_shape(*target_shape)
+        self.move_to(mobject)
+        return self
+
 
 class Square(Rectangle):
     def __init__(self, side_length: float = 2.0, **kwargs):

manimlib/mobject/mobject.py

@@ -1540,6 +1540,9 @@ class Mobject(object):
     def get_depth(self) -> float:
         return self.length_over_dim(2)
 
+    def get_shape(self) -> Tuple[float]:
+        return tuple(self.length_over_dim(dim) for dim in range(3))
+
     def get_coord(self, dim: int, direction: Vect3 = ORIGIN) -> float:
         """
         Meant to generalize get_x, get_y, get_z

manimlib/mobject/shape_matchers.py

@@ -27,13 +27,20 @@ class SurroundingRectangle(Rectangle):
         color: ManimColor = YELLOW,
         **kwargs
     ):
-        super().__init__(
-            width=mobject.get_width() + 2 * buff,
-            height=mobject.get_height() + 2 * buff,
-            color=color,
-            **kwargs
-        )
-        self.move_to(mobject)
+        super().__init__(color=color, **kwargs)
+        self.buff = buff
+        self.surround(mobject)
+
+    def surround(self, mobject, buff=None) -> Self:
+        self.mobject = mobject
+        self.buff = buff if buff is not None else self.buff
+        super().surround(mobject, self.buff)
+        return self
+
+    def set_buff(self, buff) -> Self:
+        self.buff = buff
+        self.surround(self.mobject)
+        return self
 
 
 class BackgroundRectangle(SurroundingRectangle):