diff --git a/src/zimscraperlib/image/convertion.py b/src/zimscraperlib/image/convertion.py index 31674236..3d5e0486 100644 --- a/src/zimscraperlib/image/convertion.py +++ b/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. diff --git a/src/zimscraperlib/image/transformation.py b/src/zimscraperlib/image/transformation.py index 4db56cc6..287b7efb 100644 --- a/src/zimscraperlib/image/transformation.py +++ b/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 diff --git a/src/zimscraperlib/image/utils.py b/src/zimscraperlib/image/utils.py index 2568492d..6427a298 100644 --- a/src/zimscraperlib/image/utils.py +++ b/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 diff --git a/src/zimscraperlib/zim/filesystem.py b/src/zimscraperlib/zim/filesystem.py index e0f3d6a2..0a81a7ee 100644 --- a/src/zimscraperlib/zim/filesystem.py +++ b/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 diff --git a/src/zimscraperlib/zim/items.py b/src/zimscraperlib/zim/items.py index a3d20d74..5f4c6ed0 100644 --- 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 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) diff --git a/tests/image/test_image.py b/tests/image/test_image.py index 28041737..0953a093 100644 --- a/tests/image/test_image.py +++ b/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)], diff --git a/tests/zim/test_zim_creator.py b/tests/zim/test_zim_creator.py index 9db1b1ca..a8770c4b 100644 --- a/tests/zim/test_zim_creator.py +++ b/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]), )