Add border.py and tests. Finish the logic of adding border

john-liu2 · john-liu2 · commit 4ae7c261917b · 2025-08-17T23:07:57.000-07:00
diff --git a/README.md b/README.md
@@ -42,13 +42,13 @@ Commands:
 ✗ batch_img border --help
 Usage: batch_img border [OPTIONS] SRC_PATH
 
-  Add border to image file(s)
+  Add internal border to image file(s), not expand the size
 
 Options:
   -bw, --border_width INTEGER RANGE
                                   Add border to image file(s) with the
                                   border_width. 0 - no border  [default: 5;
-                                  x>=0]
+                                  0<=x<=20]
   -bc, --border_color TEXT        Add border to image file(s) with the
                                   border_color string  [default: gray]
   -o, --output TEXT               Output file path. If skipped, use the
diff --git a/batch_img/border.py b/batch_img/border.py
@@ -0,0 +1,104 @@
+"""class Border: add border to the image file(s)
+Copyright © 2025 John Liu
+"""
+
+from pathlib import Path
+
+import piexif
+import pillow_heif
+from loguru import logger
+from PIL import Image
+
+from batch_img.common import Common
+
+pillow_heif.register_heif_opener()  # allow Pillow to open HEIC files
+
+
+class Border:
+    @staticmethod
+    def get_crop_box(width, height, border_width) -> tuple[float, float, float, float]:
+        """Get the crop box tuple
+
+        Args:
+            width: image width int
+            height: image height int
+            border_width: border width int
+
+        Returns:
+            tuple[float, float, float, float]
+        """
+        crop_left = border_width
+        crop_top = border_width
+        crop_right = width - border_width
+        crop_bottom = height - border_width
+        return crop_left, crop_top, crop_right, crop_bottom
+
+    @staticmethod
+    def add_border_1_image(
+        in_path: Path, out_path: Path, border_width: int, border_color: str
+    ) -> tuple:
+        """Add internal border to an image file, not to expand the size
+
+        Args:
+            in_path: input file path
+            out_path: output dir path
+            border_width: border width int
+            border_color: border color str
+
+        Returns:
+            tuple: bool, str
+        """
+        try:
+            with Image.open(in_path) as img:
+                width, height = img.size
+                box = Border.get_crop_box(width, height, border_width)
+                cropped_img = img.crop(box)
+                bd_img = Image.new(img.mode, (width, height), border_color)
+                bd_img.paste(cropped_img, (border_width, border_width))
+
+                out_path.mkdir(parents=True, exist_ok=True)
+                out_file = out_path
+                if out_path.is_dir():
+                    filename = f"{in_path.stem}_bw{border_width}{in_path.suffix}"
+                    out_file = Path(f"{out_path}/{filename}")
+                exif_dict = None
+                if "exif" in img.info:
+                    exif_dict = piexif.load(img.info["exif"])
+                if exif_dict:
+                    exif_bytes = piexif.dump(exif_dict)
+                    bd_img.save(out_file, img.format, optimize=True, exif=exif_bytes)
+                else:
+                    bd_img.save(out_file, img.format, optimize=True)
+            logger.info(f"Saved {out_file}")
+            return True, out_file
+        except (AttributeError, FileNotFoundError, ValueError) as e:
+            return False, f"{in_path}:\n{e}"
+
+    @staticmethod
+    def add_border_all_in_dir(
+        in_path: Path, out_path: Path, border_width: int, border_color: str
+    ) -> bool:
+        """Add border to all image files in the given dir
+
+        Args:
+            in_path: input dir path
+            out_path: output dir path
+            border_width: border width int
+            border_color: border color str
+
+        Returns:
+            bool: True - Success. False - Error
+        """
+        image_files = Common.prepare_all_files(in_path, out_path)
+        if not image_files:
+            logger.error(f"No image files at {in_path}")
+            return False
+        tasks = [(f, out_path, border_width, border_color) for f in image_files]
+        files_cnt = len(tasks)
+
+        logger.info(f"Add border to {files_cnt} image files in multiprocess ...")
+        success_cnt = Common.multiprocess_progress_bar(
+            Border.add_border_1_image, "Add border to image files", tasks
+        )
+        logger.info(f"\nSuccessfully added border to {success_cnt}/{files_cnt} files")
+        return True
diff --git a/batch_img/common.py b/batch_img/common.py
@@ -227,8 +227,7 @@ def multiprocess_progress_bar(func, desc, tasks: list) -> int:
         """
         success_cnt = 0
         files_cnt = len(tasks)
-        # Limit to 4 workers if cpu cores cnt > 4
-        workers = min(cpu_count(), 4)
+        workers = max(cpu_count(), 4)
 
         with Pool(workers) as pool:
             with tqdm(total=files_cnt, desc=desc) as pbar:
diff --git a/batch_img/interface.py b/batch_img/interface.py
@@ -18,7 +18,7 @@ def cli(ctx, version):  # pragma: no cover
             click.secho(Common.get_version())
 
 
-@cli.command(help="Add border to image file(s)")
+@cli.command(help="Add internal border to image file(s), not expand the size")
 @click.argument(
     "src_path",
     required=True,
@@ -28,7 +28,7 @@ def cli(ctx, version):  # pragma: no cover
     "--border_width",
     default=5,
     show_default=True,
-    type=click.IntRange(min=0),
+    type=click.IntRange(min=0, max=20),
     help="Add border to image file(s) with the border_width. 0 - no border",
 )
 @click.option(
diff --git a/batch_img/main.py b/batch_img/main.py
@@ -9,6 +9,7 @@
 
 from loguru import logger
 
+from batch_img.border import Border
 from batch_img.const import PKG_NAME, TS_FORMAT
 from batch_img.resize import Resize
 from batch_img.rotate import Rotate
@@ -96,8 +97,22 @@ def border(options: dict) -> bool:
         """
         logger.info(f"{json.dumps(options, indent=2)}")
         Main.init_log_file()
-        # To-do
-        return True
+        in_path = Path(options["src_path"])
+        bd_width = options.get("border_width")
+        bd_color = options.get("border_color")
+        output = options.get("output")
+        if not bd_width or bd_width == 0:
+            logger.warning(f"No add border due to bad {bd_width=}")
+            return False
+        if not output:
+            output = Path(os.getcwd())
+        else:
+            output = Path(output)
+        if in_path.is_file():
+            ok, _ = Border.add_border_1_image(in_path, output, bd_width, bd_color)
+        else:
+            ok = Border.add_border_all_in_dir(in_path, output, bd_width, bd_color)
+        return ok
 
     @staticmethod
     def default_run(options: dict) -> bool:
diff --git a/batch_img/resize.py b/batch_img/resize.py
@@ -42,7 +42,6 @@ def resize_an_image(in_path: Path, out_path: Path, length: int) -> tuple:
                 exif_dict = None
                 if "exif" in img.info:
                     exif_dict = piexif.load(img.info["exif"])
-                # Save to the same format: HEIF, JPEG, PNG, etc.
                 if exif_dict:
                     exif_bytes = piexif.dump(exif_dict)
                     img.save(out_file, img.format, optimize=True, exif=exif_bytes)
@@ -69,17 +68,10 @@ def resize_all_progress_bar(in_path: Path, out_path: Path, length: int) -> bool:
         if not image_files:
             logger.error(f"No image files at {in_path}")
             return False
-        tasks = [
-            (
-                f,
-                out_path,
-                length,
-            )
-            for f in image_files
-        ]
+        tasks = [(f, out_path, length) for f in image_files]
         files_cnt = len(tasks)
 
-        logger.info(f"Resize {files_cnt} image files with multiprocess ...")
+        logger.info(f"Resize {files_cnt} image files in multiprocess ...")
         success_cnt = Common.multiprocess_progress_bar(
             Resize.resize_an_image, "Resize image files", tasks
         )
diff --git a/batch_img/rotate.py b/batch_img/rotate.py
@@ -78,17 +78,10 @@ def rotate_all_in_dir(in_path: Path, out_path: Path, angle_cw: int) -> bool:
         if not image_files:
             logger.error(f"No image files at {in_path}")
             return False
-        tasks = [
-            (
-                f,
-                out_path,
-                angle_cw,
-            )
-            for f in image_files
-        ]
+        tasks = [(f, out_path, angle_cw) for f in image_files]
         files_cnt = len(tasks)
 
-        logger.info(f"Rotate {files_cnt} image files with multiprocess ...")
+        logger.info(f"Rotate {files_cnt} image files in multiprocess ...")
         success_cnt = Common.multiprocess_progress_bar(
             Rotate.rotate_1_image_file, "Rotate image files", tasks
         )
diff --git a/tests/test_border.py b/tests/test_border.py
@@ -0,0 +1,75 @@
+"""Test border.py
+pytest -sv tests/test_border.py
+Copyright © 2025 John Liu
+"""
+
+from os.path import dirname
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+from batch_img.border import Border
+
+
+@pytest.fixture(
+    params=[
+        (
+            Path(f"{dirname(__file__)}/data/HEIC/IMG_0070.HEIC"),
+            Path(f"{dirname(__file__)}/.out/"),
+            9,
+            "green",
+            (True, Path(f"{dirname(__file__)}/.out/IMG_0070_bw9.HEIC")),
+        ),
+        (
+            Path(f"{dirname(__file__)}/data/PNG/Checkmark.PNG"),
+            Path(f"{dirname(__file__)}/.out/"),
+            18,
+            "purple",
+            (True, Path(f"{dirname(__file__)}/.out/Checkmark_bw18.PNG")),
+        ),
+        (
+            Path(f"{dirname(__file__)}/data/JPG/152.JPG"),
+            Path(f"{dirname(__file__)}/.out/"),
+            4,
+            "#AABBCC",
+            (True, Path(f"{dirname(__file__)}/.out/152_bw4.JPG")),
+        ),
+    ]
+)
+def data_add_border_1_image(request):
+    return request.param
+
+
+def test_add_border_1_image(data_add_border_1_image):
+    in_path, out_path, bd_width, bd_color, expected = data_add_border_1_image
+    actual = Border.add_border_1_image(in_path, out_path, bd_width, bd_color)
+    assert actual == expected
+
+
+@patch("PIL.Image.open")
+def test_error_add_border_1_image(mock_open):
+    mock_open.side_effect = ValueError("VE")
+    actual = Border.add_border_1_image(Path("img/file"), Path("out/path"), 9, "red")
+    assert "img/file" in actual[1]
+
+
+@pytest.fixture(
+    params=[
+        (
+            Path(f"{dirname(__file__)}/data/mixed"),
+            Path(f"{dirname(__file__)}/.out/"),
+            9,
+            "#CCBBAA",
+            True,
+        )
+    ]
+)
+def data_add_border_all_in_dir(request):
+    return request.param
+
+
+def test_add_border_all_in_dir(data_add_border_all_in_dir):
+    in_path, out_path, bd_width, bd_color, expected = data_add_border_all_in_dir
+    actual = Border.add_border_all_in_dir(in_path, out_path, bd_width, bd_color)
+    assert actual == expected
diff --git a/tests/test_interface.py b/tests/test_interface.py
@@ -35,6 +35,21 @@ def test_border(mock_add_border, data_border):
     assert result.output == expected
 
 
+@pytest.fixture(params=["img/file -bw -9", "img/file --border_width 21"])
+def data_error_border(request):
+    return request.param
+
+
+@patch("batch_img.main.Main.border")
+def test_error_border(mock_add_border, data_error_border):
+    _input = data_error_border
+    mock_add_border.return_value = True
+    runner = CliRunner()
+    result = runner.invoke(border, args=_input.split())
+    print(result.output)
+    assert result.exception
+
+
 @pytest.fixture(
     params=[
         ("src_path --output out/dir", True, MSG_OK),
