# Source code for psd_tools.psd.image_data

    """
    Image data section structure.

    :py:class:`ImageData` corresponds to the last section of the PSD/PSB file
    where a composited image is stored. When the file does not contain layers,
    this is the only place pixels are saved.
    """
    from __future__ import absolute_import, unicode_literals

    import io
    import logging

    import attr

    from psd_tools.compression import compress, decompress
    from psd_tools.constants import Compression
    from psd_tools.psd.base import BaseElement
    from psd_tools.utils import pack, read_fmt, write_bytes, write_fmt
    from psd_tools.validators import in_

    logger = logging.getLogger(__name__)

    @attr.s(repr=False, slots=True)
    class ImageData(BaseElement):
        """
        Merged channel image data.

        .. py:attribute:: compression

            See :py:class:`~psd_tools.constants.Compression`.

        .. py:attribute:: data

            `bytes` as compressed in the `compression` flag.
        """

        compression = attr.ib(
            default=Compression.RAW, converter=Compression, validator=in_(Compression)
        )
        data = attr.ib(default=b"", type=bytes)

        @classmethod
        def read(cls, fp):
            start_pos = fp.tell()
            compression = Compression(read_fmt("H", fp)[0])
            data = fp.read()  # TODO: Parse data here. Need header.
            logger.debug("  read image data, len=%d" % (fp.tell() - start_pos))
            return cls(compression, data)

        def write(self, fp):
            start_pos = fp.tell()
            written = write_fmt(fp, "H", self.compression.value)
            written += write_bytes(fp, self.data)
            logger.debug("  wrote image data, len=%d" % (fp.tell() - start_pos))
            return written

        def get_data(self, header, split=True):
            """
            Get decompressed data.

            :param header: See :py:class:`~psd_tools.psd.header.FileHeader`.
            :return: `list` of bytes corresponding each channel.
            """
            data = decompress(
                self.data,
                self.compression,
                header.width,
                header.height * header.channels,
                header.depth,
                header.version,
            )
            if split:
                plane_size = len(data) // header.channels
                with io.BytesIO(data) as f:
                    return [f.read(plane_size) for _ in range(header.channels)]
            return data

        def set_data(self, data, header):
            """
            Set raw data and compress.

            :param data: list of raw data bytes corresponding channels.
            :param compression: compression type,
                see :py:class:`~psd_tools.constants.Compression`.
            :param header: See :py:class:`~psd_tools.psd.header.FileHeader`.
            :return: length of compressed data.
            """
            self.data = compress(
                b"".join(data),
                self.compression,
                header.width,
                header.height * header.channels,
                header.depth,
                header.version,
            )
            return len(self.data)

        @classmethod
        def new(cls, header, color=0, compression=Compression.RAW):
            """
            Create a new image data object.

            :param header: FileHeader.
            :param compression: compression type.
            :param color: default color. int or iterable for channel length.
            """
            plane_size = header.width * header.height
            if isinstance(color, (bool, int, float)):
                color = (color,) * header.channels
            if len(color) != header.channels:
                raise ValueError(
                    "Invalid color %s for channel size %d" % (color, header.channels)
                )
            # Bitmap is not supported here.
            fmt = {8: "B", 16: "H", 32: "I"}.get(header.depth)
            data = []
            for i in range(header.channels):
                data.append(pack(fmt, color[i]) * plane_size)
            self = cls(compression=compression)
            self.set_data(data, header)
            return self