libzim StringProvider support both bytes and str for content

benoit74 · benoit74 · commit d2895ee25e2b · 2024-02-14T08:48:15.000+01:00
diff --git a/src/zimscraperlib/zim/items.py b/src/zimscraperlib/zim/items.py
@@ -9,7 +9,7 @@
 import re
 import tempfile
 import urllib.parse
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import libzim.writer  # pyright: ignore
 
@@ -72,7 +72,7 @@ class StaticItem(Item):
 
     def __init__(
         self,
-        content: Optional[str] = None,
+        content: Optional[Union[str, bytes]] = None,
         fileobj: Optional[io.IOBase] = None,
         filepath: Optional[pathlib.Path] = None,
         path: Optional[str] = None,
@@ -95,6 +95,8 @@ def get_contentprovider(self) -> libzim.writer.ContentProvider:
         # content was set manually
         content = getattr(self, "content", None)
         if content is not None:
+            if not isinstance(content, (str, bytes)):
+                raise AttributeError(f"Unexpected type for content: {type(content)}")
             return StringProvider(content=content, ref=self)
 
         # using a file-like object
diff --git a/src/zimscraperlib/zim/providers.py b/src/zimscraperlib/zim/providers.py
@@ -31,7 +31,7 @@ def __init__(
 
 
 class StringProvider(libzim.writer.StringProvider):
-    def __init__(self, content: str, ref: Optional[object] = None):
+    def __init__(self, content: Union[str, bytes], ref: Optional[object] = None):
         super().__init__(content)
         self.ref = ref
 
diff --git a/tests/zim/conftest.py b/tests/zim/conftest.py
@@ -24,6 +24,26 @@ def html_str():
 """
 
 
+@pytest.fixture(scope="function")
+def html_str_cn():
+    """sample HTML content with chinese characters"""
+    return """<html>
+<body>
+<ul>
+    <li><a href="download/toto.pdf">PDF doc in 汉字</a></li>
+    <li><a href="download/toto.txt">text file</a></li>
+    <li><a href="dest.html">HTML link</a></li>
+    <li><a href="no-extension">no ext link</a></li>
+    <li><a href="http://www.example.com/index/sample.html">external link</a></li>
+    <li><a href="mailto:example@example.com">e-mail link</a></li>
+    <li><a media="">no href link</a></li>
+<object data="download/toto.jpg" width="300" height="200"></object>
+<script src="assets/js/bootstrap/bootsrap.css?v=20190101"></script>
+</body>
+</html>
+"""
+
+
 @pytest.fixture(scope="function")
 def html_file(tmp_path, html_str):
     fpath = tmp_path / "test.html"
diff --git a/tests/zim/test_zim_creator.py b/tests/zim/test_zim_creator.py
@@ -45,7 +45,7 @@ def get_contentprovider(self):
         return FileLikeProvider(self.fileobj)
 
 
-def test_zim_creator(tmp_path, png_image, html_file, html_str):
+def test_zim_creator(tmp_path, png_image, html_file, html_str: str, html_str_cn: str):
     fpath = tmp_path / "test.zim"
     main_path = "welcome"
     tags = ";".join(["toto", "tata"])
@@ -56,6 +56,13 @@ def test_zim_creator(tmp_path, png_image, html_file, html_str):
     ) as creator:
         # verbatim HTML from string
         creator.add_item_for("welcome", "wel", content=html_str, is_front=True)
+        # verbatim HTML from bytes
+        creator.add_item_for(
+            "welcome1", "wel1", content=html_str.encode(), is_front=True
+        )
+        creator.add_item_for(
+            "welcome2", "wel2", content=html_str_cn.encode("gb2312"), is_front=True
+        )
         # verbatim HTML from file
         creator.add_item_for("welcome3", "wel3", fpath=html_file)
         creator.add_item_for("welcome4", "wel4", fpath=html_file)
@@ -98,6 +105,8 @@ def test_zim_creator(tmp_path, png_image, html_file, html_str):
 
     # ensure non-rewritten articles have not been rewritten
     assert bytes(reader.get_item("welcome").content).decode(UTF8) == html_str
+    assert bytes(reader.get_item("welcome1").content).decode(UTF8) == html_str
+    assert bytes(reader.get_item("welcome2").content).decode("gb2312") == html_str_cn
     assert bytes(reader.get_item("welcome3").content).decode(UTF8) == html_str
 
     # ensure illustration is present and corrext
@@ -180,6 +189,19 @@ def test_add_item_for_delete_fail(tmp_path, png_image):
     assert reader.get_item("index")
 
 
+def test_add_item_for_unsupported_content_type(tmp_path):
+    fpath = tmp_path / "test.zim"
+    # test with incorrect content type
+    with Creator(fpath, "welcome").config_dev_metadata() as creator:
+        with pytest.raises(RuntimeError):
+            creator.add_item_for(
+                path="welcome",
+                title="hello",
+                mimetype="text/plain",
+                content=123,  # pyright: ignore[reportArgumentType]
+            )
+
+
 def test_compression(tmp_path):
     fpath = tmp_path / "test.zim"
     with Creator(
