From e317cfa47e854dd4d47ac4f2f2b65131864acb95 Mon Sep 17 00:00:00 2001
From: Evgeny Ivanov <ivanov.evgeny.n@gmail.com>
Date: Thu, 15 Jun 2023 19:12:24 +0300
Subject: [PATCH] Append ContentBase object as is (in Obj)

---
 src/rico/_content.py   | 15 ++++++++++-----
 tests/test__content.py | 38 +++++++++++++++++++++++++++-----------
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/src/rico/_content.py b/src/rico/_content.py
index ee37cef..a2417c0 100644
--- a/src/rico/_content.py
+++ b/src/rico/_content.py
@@ -356,18 +356,23 @@ def __init__(self, *objects: Any, class_: str | None = None):
         """
         super().__init__(class_=class_)
         for obj in objects:
-            if (
+            if isinstance(obj, ContentBase):
+                elements = (obj.container,)
+            elif (
                 alt is not None and isinstance(obj, alt.TopLevelMixin) or
                 plt is not None and isinstance(obj, plt.Axes | plt.Figure) or  # type: ignore  # noqa: E501
                 so is not None and isinstance(obj, so.Plot)
             ):
-                content = Chart(obj)
+                elements = Chart(obj).container
             elif hasattr(obj, "_repr_html_") and callable(obj._repr_html_):
-                content = HTML(obj._repr_html_(), strip_dataframe_borders=True)
+                elements = HTML(
+                    obj._repr_html_(),
+                    strip_dataframe_borders=True,
+                ).container
             else:
-                content = Text(obj)
+                elements = Text(obj).container
 
-            for element in content.container:
+            for element in elements:
                 self.container.append(element)
 
 
diff --git a/tests/test__content.py b/tests/test__content.py
index 60fff59..df33a22 100644
--- a/tests/test__content.py
+++ b/tests/test__content.py
@@ -424,17 +424,25 @@ class ReprHTML:
         def _repr_html_(self) -> str:
             return "<h1>Hello</h1>"
 
-    content = rico._content.Obj(ReprHTML(), "world", pyplot_axes, class_="row")
+    content_base = rico._content.ContentBase(class_="col")
 
-    div = content.container
-    assert isinstance(div, ET.Element)
-    assert div.tag == "div"
-    assert div.attrib == {"class": "row"}
-    assert div.text is None
-    assert div.tail is None
-    assert len(div) == 3
+    content = rico._content.Obj(
+        ReprHTML(),
+        "world",
+        pyplot_axes,
+        content_base,
+        class_="row",
+    )
 
-    h1 = tuple(div)[0]
+    div0 = content.container
+    assert isinstance(div0, ET.Element)
+    assert div0.tag == "div"
+    assert div0.attrib == {"class": "row"}
+    assert div0.text is None
+    assert div0.tail is None
+    assert len(div0) == 4
+
+    h1 = tuple(div0)[0]
     assert isinstance(h1, ET.Element)
     assert h1.tag == "h1"
     assert h1.attrib == {}
@@ -442,7 +450,7 @@ def _repr_html_(self) -> str:
     assert h1.tail is None
     assert len(h1) == 0
 
-    p = tuple(div)[1]
+    p = tuple(div0)[1]
     assert isinstance(p, ET.Element)
     assert p.tag == "p"
     assert p.attrib == {}
@@ -450,10 +458,18 @@ def _repr_html_(self) -> str:
     assert p.tail is None
     assert len(p) == 0
 
-    img = tuple(div)[2]
+    img = tuple(div0)[2]
     assert isinstance(img, ET.Element)
     assert img.tag == "img"
 
+    div1 = tuple(div0)[3]
+    assert isinstance(div1, ET.Element)
+    assert div1.tag == "div"
+    assert div1.attrib == {"class": "col"}
+    assert div1.text is None
+    assert div1.tail is None
+    assert len(div1) == 0
+
 
 @pytest.mark.parametrize("defer", [True, False], ids=["defer", "not defer"])
 def test_script_text(defer: bool):