shibukawa / imagesize_py

MIT License
222 stars 43 forks source link

Extended function to support Buffer and io.BufferedReader. #40

Closed TinDang97 closed 4 years ago

TinDang97 commented 4 years ago

I maintained your function to process with Bytes and io.BufferedReader. If you want, you can get it. Cause it helpful when you work with buffer.

**Note: But this way, it doesn't support XML cause ElementTree works with file.

Thanks for your repo.

def image_size(src):
    """
    Implement: https://github.com/shibukawa/imagesize_py
    Return (width, height) for a given img file content
    no requirements
    :rtype Tuple[int, int]
    """
    assert isinstance(src, (bytes, io.BufferedReader, str))
    height = -1
    width = -1
    cursor = 0

    if type(src) is str:
        src = open(src, 'rb')

    if type(src) is io.BufferedReader:
        buffer = src.read()
        src.close()
    else:
        buffer = src

    head = buffer[:24]
    size = len(head)
    # handle GIFs
    if size >= 10 and head[:6] in (b'GIF87a', b'GIF89a'):
        # Check to see if content_type is correct
        try:
            width, height = struct.unpack("<hh", head[6:10])
        except struct.error:
            raise ValueError("Invalid GIF file")
    # see png edition spec bytes are below chunk length then and finally the
    elif size >= 24 and head.startswith(b'\211PNG\r\n\032\n') and head[12:16] == b'IHDR':
        try:
            width, height = struct.unpack(">LL", head[16:24])
        except struct.error:
            raise ValueError("Invalid PNG file")
    # Maybe this is for an older PNG version.
    elif size >= 16 and head.startswith(b'\211PNG\r\n\032\n'):
        # Check to see if we have the right content type
        try:
            width, height = struct.unpack(">LL", head[8:16])
        except struct.error:
            raise ValueError("Invalid PNG file")
    # handle JPEGs
    elif size >= 2 and head.startswith(b'\377\330'):
        try:
            size = 2
            ftype = 0
            while not 0xc0 <= ftype <= 0xcf or ftype in [0xc4, 0xc8, 0xcc]:
                cursor += size
                byte = buffer[cursor:cursor+1]
                cursor += 1
                while ord(byte) == 0xff:
                    byte = buffer[cursor:cursor+1]
                    cursor += 1
                ftype = ord(byte)
                size = struct.unpack('>H', buffer[cursor:cursor+2])[0] - 2
                cursor += 2
            # We are at a SOFn block
            cursor += 1  # Skip `precision' byte.
            height, width = struct.unpack('>HH', buffer[cursor:cursor+4])
            cursor += 4
        except struct.error:
            raise ValueError("Invalid JPEG file")
    # handle JPEG2000s
    elif size >= 12 and head.startswith(b'\x00\x00\x00\x0cjP  \r\n\x87\n'):
        cursor = 48
        try:
            height, width = struct.unpack('>LL', buffer[cursor:cursor+8])
        except struct.error:
            raise ValueError("Invalid JPEG2000 file")
    # handle big endian TIFF
    elif size >= 8 and head.startswith(b"\x4d\x4d\x00\x2a"):
        offset = struct.unpack('>L', head[4:8])[0]
        cursor = offset
        ifdsize = struct.unpack(">H", buffer[cursor:cursor+2])[0]
        cursor += 2
        for i in range(ifdsize):
            tag, datatype, count, data = struct.unpack(">HHLL", buffer[cursor:cursor+12])
            if tag == 256:
                if datatype == 3:
                    width = int(data / 65536)
                elif datatype == 4:
                    width = data
                else:
                    raise ValueError("Invalid TIFF file: width column data type should be SHORT/LONG.")
            elif tag == 257:
                if datatype == 3:
                    height = int(data / 65536)
                elif datatype == 4:
                    height = data
                else:
                    raise ValueError("Invalid TIFF file: height column data type should be SHORT/LONG.")
            if width != -1 and height != -1:
                break
        if width == -1 or height == -1:
            raise ValueError("Invalid TIFF file: width and/or height IDS entries are missing.")
    elif size >= 8 and head.startswith(b"\x49\x49\x2a\x00"):
        offset = struct.unpack('<L', head[4:8])[0]
        cursor = offset
        ifdsize = struct.unpack("<H", buffer[cursor:cursor+2])[0]
        cursor += 2
        for i in range(ifdsize):
            tag, datatype, count, data = struct.unpack("<HHLL", buffer[cursor:cursor+12])
            if tag == 256:
                width = data
            elif tag == 257:
                height = data
            if width != -1 and height != -1:
                break
        if width == -1 or height == -1:
            raise ValueError("Invalid TIFF file: width and/or height IDS entries are missing.")
    return width, height