Skip to content

Commit

Permalink
STY: Refactor _xobj_to_image
Browse files Browse the repository at this point in the history
* Create _apply_decode and move it to _xobj_image_helpers
* Create _get_mode_and_invert_color and move it to _xobj_image_helpers
* Create _apply_alpha as an inner function as it's calling _xobj_to_image

This reduced the Cyclomatic Complexity from 44 to 19
  • Loading branch information
MartinThoma committed Sep 22, 2024
1 parent 966e015 commit eeef7cc
Show file tree
Hide file tree
Showing 3 changed files with 129 additions and 86 deletions.
74 changes: 73 additions & 1 deletion pypdf/_xobj_image_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@

import sys
from io import BytesIO
from typing import Any, List, Literal, Tuple, Union, cast
from typing import Any, Dict, List, Literal, Tuple, Union, cast

from ._utils import check_if_whitespace_only, logger_warning
from .constants import ColorSpaces
from .constants import FilterTypes as FT
from .constants import ImageAttributes as IA
from .errors import EmptyImageDataError, PdfReadError
from .generic import (
ArrayObject,
Expand Down Expand Up @@ -303,3 +305,73 @@ def _handle_jpx(
img = img.convert("RGB")
image_format = "JPEG2000"
return img, image_format, extension, invert_color


def _apply_decode(
img: Image.Image,
x_object_obj: Dict[str, Any],
lfilters: FT,
color_space: Union[str, List[Any], Any],
invert_color: bool,
) -> Image.Image:
# CMYK image and other colorspaces without decode
# requires reverting scale (cf p243,2§ last sentence)
decode = x_object_obj.get(
IA.DECODE,
([1.0, 0.0] * len(img.getbands()))
if (
(img.mode == "CMYK" and lfilters in (FT.DCT_DECODE, FT.JPX_DECODE))
or (invert_color and img.mode == "L")
)
else None,
)
if (
isinstance(color_space, ArrayObject)
and color_space[0].get_object() == "/Indexed"
):
decode = None # decode is meanless of Indexed
if (
isinstance(color_space, ArrayObject)
and color_space[0].get_object() == "/Separation"
):
decode = [1.0, 0.0] * len(img.getbands())
if decode is not None and not all(decode[i] == i % 2 for i in range(len(decode))):
lut: List[int] = []
for i in range(0, len(decode), 2):
dmin = decode[i]
dmax = decode[i + 1]
lut.extend(
round(255.0 * (j / 255.0 * (dmax - dmin) + dmin)) for j in range(256)
)
img = img.point(lut)
return img


def _get_mode_and_invert_color(
x_object_obj: Dict[str, Any], colors: int, color_space: Union[str, List[Any], Any]
) -> Tuple[mode_str_type, bool]:
if (
IA.COLOR_SPACE in x_object_obj
and x_object_obj[IA.COLOR_SPACE] == ColorSpaces.DEVICE_RGB
):
# https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes
mode: mode_str_type = "RGB"
if x_object_obj.get("/BitsPerComponent", 8) < 8:
mode, invert_color = _get_imagemode(
f"{x_object_obj.get('/BitsPerComponent', 8)}bit", 0, ""
)
else:
mode, invert_color = _get_imagemode(
color_space,
2
if (
colors == 1
and (
not isinstance(color_space, NullObject)
and "Gray" not in color_space
)
)
else colors,
"",
)
return mode, invert_color
7 changes: 4 additions & 3 deletions pypdf/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
ISO 32000-2:2020 (PDF 2.0)
"""

from enum import IntFlag, auto
from enum import Enum, IntFlag, auto, unique
from typing import Dict, Tuple


Expand Down Expand Up @@ -168,7 +168,7 @@ class PagesAttributes:
PARENT = "/Parent" # dictionary, required; indirect reference to pages object
KIDS = "/Kids" # array, required; List of indirect references
COUNT = "/Count" # integer, required; the number of leaf nodes (page objects)
# that are descendants of this node within the page tree
# that are descendants of this node within the page tree


class PageAttributes:
Expand Down Expand Up @@ -240,7 +240,8 @@ class StreamAttributes:
DECODE_PARMS = "/DecodeParms" # variable, optional -- 'decodeParams is wrong


class FilterTypes:
@unique
class FilterTypes(Enum):
"""§7.4 of the 1.7 and 2.0 references."""

ASCII_HEX_DECODE = "/ASCIIHexDecode" # abbreviation: AHx
Expand Down
134 changes: 52 additions & 82 deletions pypdf/filters.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,6 @@
ord_,
)
from .constants import CcittFaxDecodeParameters as CCITT
from .constants import ColorSpaces
from .constants import FilterTypeAbbreviations as FTA
from .constants import FilterTypes as FT
from .constants import ImageAttributes as IA
Expand Down Expand Up @@ -735,57 +734,72 @@ def _xobj_to_image(x_object_obj: Dict[str, Any]) -> Tuple[Optional[str], bytes,
from ._xobj_image_helpers import (
Image,
UnidentifiedImageError,
_apply_decode,
_extended_image_frombytes,
_get_imagemode,
_get_mode_and_invert_color,
_handle_flate,
_handle_jpx,
mode_str_type,
)

def _apply_alpha(
img: Image.Image,
x_object_obj: Dict[str, Any],
obj_as_text: str,
image_format: str,
extension: str,
) -> Tuple[Image.Image, str, str]:
alpha = None
if IA.S_MASK in x_object_obj: # add alpha channel
alpha = _xobj_to_image(x_object_obj[IA.S_MASK])[2]
if img.size != alpha.size:
logger_warning(
f"image and mask size not matching: {obj_as_text}", __name__
)
else:
# TODO : implement mask
if alpha.mode != "L":
alpha = alpha.convert("L")
if img.mode == "P":
img = img.convert("RGB")
elif img.mode == "1":
img = img.convert("L")
img.putalpha(alpha)
if "JPEG" in image_format:
extension = ".jp2"
image_format = "JPEG2000"
else:
extension = ".png"
image_format = "PNG"
return img, extension, image_format

# for error reporting
if x_object_obj is None: # pragma: no cover
obj_as_text = x_object_obj.indirect_reference.__repr__()
else:
obj_as_text = x_object_obj.__repr__()
obj_as_text = (
x_object_obj.indirect_reference.__repr__() # type: ignore
if x_object_obj is None # pragma: no cover
else x_object_obj.__repr__()
)

# Extract size and data
size = (cast(int, x_object_obj[IA.WIDTH]), cast(int, x_object_obj[IA.HEIGHT]))
data = x_object_obj.get_data() # type: ignore
if isinstance(data, str): # pragma: no cover
data = data.encode()
if len(data) % (size[0] * size[1]) == 1 and data[-1] == 0x0A: # ie. '\n'
data = data[:-1]

# Determine color properties
colors = x_object_obj.get("/Colors", 1)
color_space: Any = x_object_obj.get("/ColorSpace", NullObject()).get_object()
if isinstance(color_space, list) and len(color_space) == 1:
color_space = color_space[0].get_object()
if (
IA.COLOR_SPACE in x_object_obj
and x_object_obj[IA.COLOR_SPACE] == ColorSpaces.DEVICE_RGB
):
# https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes
mode: mode_str_type = "RGB"
if x_object_obj.get("/BitsPerComponent", 8) < 8:
mode, invert_color = _get_imagemode(
f"{x_object_obj.get('/BitsPerComponent', 8)}bit", 0, ""
)
else:
mode, invert_color = _get_imagemode(
color_space,
2
if (
colors == 1
and (
not isinstance(color_space, NullObject)
and "Gray" not in color_space
)
)
else colors,
"",
)
extension = None
alpha = None

mode, invert_color = _get_mode_and_invert_color(x_object_obj, colors, color_space)

# Determine filters
filters = x_object_obj.get(SA.FILTER, NullObject()).get_object()
lfilters = filters[-1] if isinstance(filters, list) else filters

extension = None
if lfilters in (FT.FLATE_DECODE, FT.RUN_LENGTH_DECODE):
img, image_format, extension, _ = _handle_flate(
size,
Expand Down Expand Up @@ -839,57 +853,13 @@ def _xobj_to_image(x_object_obj: Dict[str, Any]) -> Tuple[Optional[str], bytes,
".png",
False,
)
# CMYK image and other colorspaces without decode
# requires reverting scale (cf p243,2§ last sentence)
decode = x_object_obj.get(
IA.DECODE,
([1.0, 0.0] * len(img.getbands()))
if (
(img.mode == "CMYK" and lfilters in (FT.DCT_DECODE, FT.JPX_DECODE))
or (invert_color and img.mode == "L")
)
else None,
)
if (
isinstance(color_space, ArrayObject)
and color_space[0].get_object() == "/Indexed"
):
decode = None # decode is meanless of Indexed
if (
isinstance(color_space, ArrayObject)
and color_space[0].get_object() == "/Separation"
):
decode = [1.0, 0.0] * len(img.getbands())
if decode is not None and not all(decode[i] == i % 2 for i in range(len(decode))):
lut: List[int] = []
for i in range(0, len(decode), 2):
dmin = decode[i]
dmax = decode[i + 1]
lut.extend(
round(255.0 * (j / 255.0 * (dmax - dmin) + dmin)) for j in range(256)
)
img = img.point(lut)

if IA.S_MASK in x_object_obj: # add alpha channel
alpha = _xobj_to_image(x_object_obj[IA.S_MASK])[2]
if img.size != alpha.size:
logger_warning(f"image and mask size not matching: {obj_as_text}", __name__)
else:
# TODO : implement mask
if alpha.mode != "L":
alpha = alpha.convert("L")
if img.mode == "P":
img = img.convert("RGB")
elif img.mode == "1":
img = img.convert("L")
img.putalpha(alpha)
if "JPEG" in image_format:
extension = ".jp2"
image_format = "JPEG2000"
else:
extension = ".png"
image_format = "PNG"
img = _apply_decode(img, x_object_obj, lfilters, color_space, invert_color)
img, extension, image_format = _apply_alpha(
img, x_object_obj, obj_as_text, image_format, extension
)

# Save image to bytes
img_byte_arr = BytesIO()
try:
img.save(img_byte_arr, format=image_format)
Expand Down

0 comments on commit eeef7cc

Please sign in to comment.