STY: Slight refactor of remove_images (#2979)

Increase readability.
py-pdf · Dec 3, 2024 · 273732b · 273732b
1 parent fa1838c
commit 273732b
Showing 1 changed file with 8 additions and 18 deletions.
diff --git a/pypdf/_writer.py b/pypdf/_writer.py
@@ -2194,29 +2194,19 @@ def remove_images(
         Remove images from this output.
 
         Args:
-            to_delete : The type of images to be deleted
+            to_delete: The type of images to be deleted
                 (default = all images types)
 
         """
         if isinstance(to_delete, bool):
             to_delete = ImageType.ALL
-        i = (
-            (
-                ObjectDeletionFlag.XOBJECT_IMAGES
-                if to_delete & ImageType.XOBJECT_IMAGES
-                else ObjectDeletionFlag.NONE
-            )
-            | (
-                ObjectDeletionFlag.INLINE_IMAGES
-                if to_delete & ImageType.INLINE_IMAGES
-                else ObjectDeletionFlag.NONE
-            )
-            | (
-                ObjectDeletionFlag.DRAWING_IMAGES
-                if to_delete & ImageType.DRAWING_IMAGES
-                else ObjectDeletionFlag.NONE
-            )
-        )
+
+        i = ObjectDeletionFlag.NONE
+
+        for image in ("XOBJECT_IMAGES", "INLINE_IMAGES", "DRAWING_IMAGES"):
+            if to_delete & ImageType[image]:
+                i |= ObjectDeletionFlag[image]
+
         for page in self.pages:
             self.remove_objects_from_page(page, i)