Merge pull request #132 from openzim/fix_type_hints

Fix type hints

Author: benoit74

src/zimscraperlib/image/convertion.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 # vim: ai ts=4 sts=4 et sw=4 nu
 
+import io
 import pathlib
-from typing import Optional
+from typing import Union
 
 import PIL
 
@@ -13,7 +14,9 @@
 
 
 def convert_image(
-    src: pathlib.Path, dst: pathlib.Path, **params: Optional[dict]
+    src: Union[pathlib.Path, io.BytesIO],
+    dst: Union[pathlib.Path, io.BytesIO],
+    **params: str,
 ) -> None:
     """convert an image file from one format to another
     params: Image.save() parameters. Depends on dest format.

src/zimscraperlib/image/transformation.py

@@ -19,7 +19,7 @@ def resize_image(
     dst: Optional[Union[pathlib.Path, io.BytesIO]] = None,
     method: Optional[str] = "width",
     allow_upscaling: Optional[bool] = True,  # noqa: FBT002
-    **params: Optional[dict],
+    **params: str,
 ) -> None:
     """resize an image to requested dimensions
 

src/zimscraperlib/image/utils.py

@@ -1,17 +1,18 @@
 #!/usr/bin/env python
 # vim: ai ts=4 sts=4 et sw=4 nu
 
+import io
 import pathlib
-from typing import Optional
+from typing import Optional, Union
 
 from PIL import Image
 
 
 def save_image(
     src: Image,  # pyright: ignore
-    dst: pathlib.Path,
+    dst: Union[pathlib.Path, io.BytesIO],
     fmt: Optional[str] = None,
-    **params: Optional[dict],
+    **params: str,
 ) -> None:
     """PIL.Image.save() wrapper setting default parameters"""
     args = {"JPEG": {"quality": 100}, "PNG": {}}.get(fmt, {})  # pyright: ignore

src/zimscraperlib/zim/filesystem.py

@@ -46,8 +46,10 @@ def __init__(
         self,
         root: pathlib.Path,
         filepath: pathlib.Path,
-    ):  # pyright: ignore
-        super().__init__(root=root, filepath=filepath)  # pyright: ignore
+    ):
+        super().__init__()
+        self.root = root
+        self.filepath = filepath
         # first look inside the file's magic headers
         self.mimetype = get_file_mimetype(self.filepath)
         # most web-specific files are plain text. In this case, use extension

src/zimscraperlib/zim/items.py

@@ -9,7 +9,7 @@
 import re
 import tempfile
 import urllib.parse
-from typing import Dict, Union
+from typing import Any, Optional
 
 import libzim.writer  # pyright: ignore
 
@@ -23,12 +23,25 @@
 
 
 class Item(libzim.writer.Item):
-    """libzim.writer.Item returning props for path/title/mimetype plus a callback
-
-    Calls your `callback` prop on deletion"""
-
-    def __init__(self, **kwargs: Dict[str, Union[str, bool, bytes]]):
+    """libzim.writer.Item returning props for path/title/mimetype"""
+
+    def __init__(
+        self,
+        path: Optional[str] = None,
+        title: Optional[str] = None,
+        mimetype: Optional[str] = None,
+        hints: Optional[dict] = None,
+        **kwargs: Any,
+    ):
         super().__init__()
+        if path:
+            kwargs["path"] = path
+        if title:
+            kwargs["title"] = title
+        if mimetype:
+            kwargs["mimetype"] = mimetype
+        if hints:
+            kwargs["hints"] = hints
         for k, v in kwargs.items():
             setattr(self, k, v)
 
@@ -57,21 +70,45 @@ class StaticItem(Item):
     more efficiently: now when the libzim destroys the CP, python will destroy
     the Item and we can be notified that we're effectively through with our content"""
 
+    def __init__(
+        self,
+        content: Optional[str] = None,
+        fileobj: Optional[io.IOBase] = None,
+        filepath: Optional[pathlib.Path] = None,
+        path: Optional[str] = None,
+        title: Optional[str] = None,
+        mimetype: Optional[str] = None,
+        hints: Optional[dict] = None,
+        **kwargs: Any,
+    ):
+        if content:
+            kwargs["content"] = content
+        if fileobj:
+            kwargs["fileobj"] = fileobj
+        if filepath:
+            kwargs["filepath"] = filepath
+        super().__init__(
+            path=path, title=title, mimetype=mimetype, hints=hints, **kwargs
+        )
+
     def get_contentprovider(self) -> libzim.writer.ContentProvider:
         # content was set manually
-        if getattr(self, "content", None) is not None:
-            return StringProvider(content=self.content, ref=self)
+        content = getattr(self, "content", None)
+        if content is not None:
+            return StringProvider(content=content, ref=self)
 
         # using a file-like object
-        if getattr(self, "fileobj", None):
+        fileobj = getattr(self, "fileobj", None)
+        if fileobj:
             return FileLikeProvider(
-                fileobj=self.fileobj, ref=self, size=getattr(self, "size", None)
+                fileobj=fileobj, ref=self, size=getattr(self, "size", None)
             )
 
         # we had to download locally to get size
-        if getattr(self, "filepath", None):
+        filepath = getattr(self, "filepath", None)
+        if filepath:
             return FileProvider(
-                filepath=self.filepath, ref=self, size=getattr(self, "size", None)
+                filepath=filepath, ref=self, size=getattr(self, "size", None)
             )
 
         raise NotImplementedError("No data to provide`")
@@ -106,8 +143,21 @@ def download_for_size(url, on_disk, tmp_dir=None):
         size, _ = stream_file(url.geturl(), fpath=fpath, byte_stream=stream)
         return fpath or stream, size
 
-    def __init__(self, url: str, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(
+        self,
+        url: str,
+        path: Optional[str] = None,
+        title: Optional[str] = None,
+        mimetype: Optional[str] = None,
+        hints: Optional[dict] = None,
+        use_disk: Optional[bool] = None,
+        **kwargs: Any,
+    ):
+        if use_disk:
+            kwargs["use_disk"] = use_disk
+        super().__init__(
+            path=path, title=title, mimetype=mimetype, hints=hints, **kwargs
+        )
         self.url = urllib.parse.urlparse(url)
         use_disk = getattr(self, "use_disk", False)
 

tests/image/test_image.py

@@ -296,6 +296,30 @@ def test_change_image_format_defaults(png_image, tmp_path):
     assert dst_image.format == "WEBP"
 
 
+def test_convert_io_src_dst(png_image: pathlib.Path):
+    src = io.BytesIO(png_image.read_bytes())
+    dst = io.BytesIO()
+    convert_image(src, dst, fmt="PNG")
+    dst_image = Image.open(dst)
+    assert dst_image.format == "PNG"
+
+
+def test_convert_io_src_path_dst(png_image: pathlib.Path, tmp_path: pathlib.Path):
+    src = io.BytesIO(png_image.read_bytes())
+    dst = tmp_path / "test.png"
+    convert_image(src, dst, fmt="PNG")
+    dst_image = Image.open(dst)
+    assert dst_image.format == "PNG"
+
+
+def test_convert_path_src_io_dst(png_image: pathlib.Path):
+    src = png_image
+    dst = io.BytesIO()
+    convert_image(src, dst, fmt="PNG")
+    dst_image = Image.open(dst)
+    assert dst_image.format == "PNG"
+
+
 @pytest.mark.parametrize(
     "fmt,exp_size",
     [("png", 128), ("jpg", 128)],

tests/zim/test_zim_creator.py

@@ -4,7 +4,6 @@
 import base64
 import datetime
 import io
-import os
 import pathlib
 import random
 import shutil
@@ -41,6 +40,8 @@ def get_contentprovider(self):
 
 class FileLikeProviderItem(StaticItem):
     def get_contentprovider(self):
+        if not self.fileobj:
+            raise AttributeError("fileobj cannot be None")
         return FileLikeProvider(self.fileobj)
 
 
@@ -119,7 +120,7 @@ def test_noindexlanguage(tmp_path):
     creator = Creator(fpath, "welcome").config_dev_metadata(Language="bam")
     creator.config_indexing(False)
     with creator as creator:
-        creator.add_item(StaticItem(path="welcome", content="hello"))  # pyright: ignore
+        creator.add_item(StaticItem(path="welcome", content="hello"))
         creator.add_item_for("index", "Index", content="-", mimetype="text/html")
 
     reader = Archive(fpath)
@@ -165,15 +166,11 @@ def test_add_item_for_delete_fail(tmp_path, png_image):
     # copy file to local path
     shutil.copyfile(png_image, local_path)
 
-    def remove_source(item):
-        os.remove(item.filepath)
-
     with Creator(fpath, "welcome").config_dev_metadata() as creator:
         creator.add_item(
             StaticItem(
-                filepath=local_path,  # pyright: ignore
-                path="index",  # pyright: ignore
-                callback=remove_source,  # pyright: ignore
+                filepath=local_path,
+                path="index",
             ),
             callback=(delete_callback, local_path),
         )
@@ -188,18 +185,18 @@ def test_compression(tmp_path):
     with Creator(
         tmp_path / "test.zim", "welcome", compression="zstd"
     ).config_dev_metadata() as creator:
-        creator.add_item(StaticItem(path="welcome", content="hello"))  # pyright: ignore
+        creator.add_item(StaticItem(path="welcome", content="hello"))
 
     with Creator(
         fpath, "welcome", compression=Compression.zstd  # pyright: ignore
     ).config_dev_metadata() as creator:
-        creator.add_item(StaticItem(path="welcome", content="hello"))  # pyright: ignore
+        creator.add_item(StaticItem(path="welcome", content="hello"))
 
 
 def test_double_finish(tmp_path):
     fpath = tmp_path / "test.zim"
     with Creator(fpath, "welcome").config_dev_metadata() as creator:
-        creator.add_item(StaticItem(path="welcome", content="hello"))  # pyright: ignore
+        creator.add_item(StaticItem(path="welcome", content="hello"))
 
     # ensure we can finish an already finished creator
     creator.finish()
@@ -219,11 +216,7 @@ def test_sourcefile_removal(tmp_path, html_file):
         # copy html to folder
         src_path = pathlib.Path(tmpdir.name, "source.html")
         shutil.copyfile(html_file, src_path)
-        creator.add_item(
-            StaticItem(
-                filepath=src_path, path=src_path.name, ref=tmpdir  # pyright: ignore
-            )
-        )
+        creator.add_item(StaticItem(filepath=src_path, path=src_path.name, ref=tmpdir))
         del tmpdir
 
     assert not src_path.exists()
@@ -241,7 +234,7 @@ def test_sourcefile_removal_std(tmp_path, html_file):
                 StaticItem(
                     filepath=paths[-1],
                     path=paths[-1].name,
-                    mimetype="text/html",  # pyright: ignore
+                    mimetype="text/html",
                 ),
                 callback=(delete_callback, paths[-1]),
             )