LinearCollection basic functionality and tests

Author: Nikhil-42

Diff for: src/manim_data_structures/__init__.py

@@ -3,6 +3,7 @@
 from .m_array import *
 from .m_enum import *
 from .m_variable import *
+from .mlinearcollection import *
 
 __all__ = [
     "MArrayElement",
@@ -12,4 +13,5 @@
     "MArrayDirection",
     "MArrayElementComp",
     "MVariable",
+    "LinearCollection",
 ]

Diff for: src/manim_data_structures/mlinearcollection.py

@@ -0,0 +1,78 @@
+import copy
+from typing import Iterable, List
+
+import numpy as np
+from manim import RIGHT, Integer, Mobject, Square, VMobject
+
+
+class LinearCollection(VMobject):
+    __data: List[Mobject]
+    __data_template: Mobject
+    __container_template: Mobject
+    __delimiter: Mobject
+    __direction: np.ndarray
+
+    def __init__(
+        self,
+        data: Iterable,
+        data_template=Integer,
+        container_template=Square,
+        delimiter=None,
+        direction=RIGHT,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.__data = []
+        self.__data_template = data_template
+        self.__container_template = container_template
+        self.__delimiter = delimiter
+        self.__direction = direction
+
+        for value in data:
+            self.__data.append(data_template(value))
+            new_container = container_template()
+            new_container.add(self.__data[-1])
+            self.add(new_container)
+
+        self.arrange(direction=self.__direction)
+
+    def insert(self, value):
+        self.__data.append(self.__data_template(value))
+        new_container = self.__container_template()
+        new_container.add(self.__data[-1])
+        self.add(new_container)
+        self.arrange(direction=self.__direction)
+
+    def remove(self):
+        pass
+
+    def __getitem__(self, key):
+        return self.__data.__getitem__(key).get_value()
+
+    def __setitem__(self, key, value):
+        return self.__data.__getitem__(key).set_value(value)
+
+    def __iter__(self):
+        # Initialize a Data iterator when we start iterating
+        self.__iter__data = iter(self.__data)
+        return self
+
+    def __next__(self):
+        try:
+            # Return the value of the next Data in the iterator
+            return next(self.__iter__data).get_value()
+        except StopIteration as e:
+            # Clean up the iterator when we're done
+            del self.__iter__data
+            raise e
+
+    def __add__(self, rhs):
+        self_copy = copy.copy(self)
+        for container in rhs.submobjects:
+            self_copy.add(copy.copy(container))
+        for data in rhs._LinearCollection__data:
+            self_copy.__data.append(copy.copy(data))
+
+        return self_copy

Diff for: tests/test_mlinearcollection.py

@@ -1,11 +1,71 @@
-# Fill with appropriate tests
-def test_getitem():
+from manim import Integer, Mobject, Square
+
+from manim_data_structures import LinearCollection
+
+
+def test_constructor():
+    data = [i + 1 for i in range(10)]
+    lc = LinearCollection(data)
+    assert type(lc) == LinearCollection
+    assert type(lc._LinearCollection__data) == list
+    assert isinstance(lc, Mobject)
+    assert isinstance(lc._LinearCollection__data[0], Mobject)
+    assert isinstance(lc.submobjects[0], Mobject)
+    assert type(lc._LinearCollection__data[0]) == Integer
+    assert type(lc.submobjects[0]) == Square
+
+    for i in range(len(data)):
+        assert lc._LinearCollection__data[i].get_value() == i + 1
+        assert lc.submobjects[i].submobjects[0].get_value() == i + 1
+
+
+def test_insertion():
     pass
 
 
-def test_setitem():
+def test_removal():
     pass
 
 
+def test_getitem():
+    data = [i + 1 for i in range(10)]
+    lc = LinearCollection(data)
+
+    for i in range(len(data)):
+        assert lc[i] == data[i]
+
+
+def test_setitem():
+    data = [i + 1 for i in range(10)]
+    lc = LinearCollection(data)
+
+    for i in range(len(data)):
+        assert lc[i] == data[i]
+        lc[i] = data[-(1 + i)]
+
+    for i in range(len(data)):
+        assert lc[i] == data[-(1 + i)]
+
+
 def test_iteration():
-    pass
+    data = [i + 1 for i in range(10)]
+    lc = LinearCollection(data)
+
+    for lc_value, data_value in zip(lc, data):
+        assert lc_value == data_value
+
+
+def test_add():
+    data = [i + 1 for i in range(10)]
+    lc1 = LinearCollection(data)
+    lc2 = LinearCollection(data[::-1])
+
+    for i in range(len(data)):
+        assert lc1[i] == data[i]
+        assert lc2[i] == data[-(1 + i)]
+
+    for i, value in enumerate(lc1 + lc2):
+        if i < len(data):
+            assert value == lc1[i]
+        else:
+            assert value == lc2[i - len(data)]