1"""Image loaders."""
2from .common import SDLError
3from .compat import UnsupportedError, byteify
4from .. import endian, surface, pixels
5
6_HASPIL = True
7try:
8    from PIL import Image
9except ImportError:
10    _HASPIL = False
11
12_HASSDLIMAGE = True
13try:
14    from .. import sdlimage
15except ImportError:
16    _HASSDLIMAGE = False
17
18__all__ = ["get_image_formats", "load_image"]
19
20
21def get_image_formats():
22    """Gets the formats supported in the default installation."""
23    if not _HASPIL and not _HASSDLIMAGE:
24        return ("bmp", )
25    return ("bmp", "cur", "gif", "ico", "jpg", "lbm", "pbm", "pcx", "pgm",
26            "png", "pnm", "ppm", "svg", "tga", "tif", "webp", "xcf", "xpm")
27
28
29def load_image(fname, enforce=None):
30    """Creates a SDL_Surface from an image file.
31
32    This function makes use of the Python Imaging Library, if it is available
33    on the target execution environment. The function will try to load the
34    file via sdl2 first. If the file could not be loaded, it will try
35    to load it via sdl2.sdlimage and PIL.
36
37    You can force the function to use only one of them, by passing the enforce
38    as either "PIL" or "SDL".
39
40    Note: This will call sdl2.sdlimage.init() implicitly with the default
41    arguments, if the module is available and if sdl2.SDL_LoadBMP() failed to
42    load the image.
43    """
44    if enforce is not None and enforce not in ("PIL", "SDL"):
45        raise ValueError("enforce must be either 'PIL' or 'SDL', if set")
46    if fname is None:
47        raise ValueError("fname must be a string")
48
49    name = fname
50    if hasattr(fname, 'encode'):
51        name = byteify(fname, "utf-8")
52
53    if not _HASPIL and not _HASSDLIMAGE:
54        imgsurface = surface.SDL_LoadBMP(name)
55        if not imgsurface:
56            raise UnsupportedError(load_image,
57                                   "cannot use PIL or SDL for image loading")
58        return imgsurface.contents
59    if enforce == "PIL" and not _HASPIL:
60        raise UnsupportedError(load_image, "cannot use PIL (not found)")
61    if enforce == "SDL" and not _HASSDLIMAGE:
62        imgsurface = surface.SDL_LoadBMP(name)
63        if not imgsurface:
64            raise UnsupportedError(load_image,
65                                   "cannot use SDL_image (not found)")
66        return imgsurface.contents
67
68    imgsurface = None
69    if enforce != "PIL" and _HASSDLIMAGE:
70        sdlimage.IMG_Init(sdlimage.IMG_INIT_JPG | sdlimage.IMG_INIT_PNG |
71                          sdlimage.IMG_INIT_TIF | sdlimage.IMG_INIT_WEBP)
72        imgsurface = sdlimage.IMG_Load(name)
73        if not imgsurface:
74            # An error occured - if we do not try PIL, break out now
75            if not _HASPIL or enforce == "SDL":
76                raise SDLError(sdlimage.IMG_GetError())
77        else:
78            imgsurface = imgsurface.contents
79
80    if enforce != "SDL" and _HASPIL and not imgsurface:
81        image = Image.open(fname)
82        mode = image.mode
83        width, height = image.size
84        rmask = gmask = bmask = amask = 0
85        if mode in ("1", "L", "P"):
86            # 1 = B/W, 1 bit per byte
87            # "L" = greyscale, 8-bit
88            # "P" = palette-based, 8-bit
89            pitch = width
90            depth = 8
91        elif mode == "RGB":
92            # 3x8-bit, 24bpp
93            if endian.SDL_BYTEORDER == endian.SDL_LIL_ENDIAN:
94                rmask = 0x0000FF
95                gmask = 0x00FF00
96                bmask = 0xFF0000
97            else:
98                rmask = 0xFF0000
99                gmask = 0x00FF00
100                bmask = 0x0000FF
101            depth = 24
102            pitch = width * 3
103        elif mode in ("RGBA", "RGBX"):
104            # RGBX: 4x8-bit, no alpha
105            # RGBA: 4x8-bit, alpha
106            if endian.SDL_BYTEORDER == endian.SDL_LIL_ENDIAN:
107                rmask = 0x000000FF
108                gmask = 0x0000FF00
109                bmask = 0x00FF0000
110                if mode == "RGBA":
111                    amask = 0xFF000000
112            else:
113                rmask = 0xFF000000
114                gmask = 0x00FF0000
115                bmask = 0x0000FF00
116                if mode == "RGBA":
117                    amask = 0x000000FF
118            depth = 32
119            pitch = width * 4
120        else:
121            # We do not support CMYK or YCbCr for now
122            raise TypeError("unsupported image format")
123
124        pxbuf = image.tobytes()
125        imgsurface = surface.SDL_CreateRGBSurfaceFrom(pxbuf, width, height,
126                                                      depth, pitch, rmask,
127                                                      gmask, bmask, amask)
128        if not imgsurface:
129            raise SDLError()
130        imgsurface = imgsurface.contents
131        # the pixel buffer must not be freed for the lifetime of the surface
132        imgsurface._pxbuf = pxbuf
133
134        if mode == "P":
135            # Create a SDL_Palette for the SDL_Surface
136            def _chunk(seq, size):
137                for x in range(0, len(seq), size):
138                    yield seq[x:x + size]
139
140            rgbcolors = image.getpalette()
141            sdlpalette = pixels.SDL_AllocPalette(len(rgbcolors) // 3)
142            if not sdlpalette:
143                raise SDLError()
144            SDL_Color = pixels.SDL_Color
145            for idx, (r, g, b) in enumerate(_chunk(rgbcolors, 3)):
146                sdlpalette.contents.colors[idx] = SDL_Color(r, g, b)
147            ret = surface.SDL_SetSurfacePalette(imgsurface, sdlpalette)
148            # This will decrease the refcount on the palette, so it gets
149            # freed properly on releasing the SDL_Surface.
150            pixels.SDL_FreePalette(sdlpalette)
151            if ret != 0:
152                raise SDLError()
153
154            # If the image has a single transparent palette index, set
155            # that index as the color key to make blitting correct.
156            if 'transparency' in image.info and isinstance(image.info['transparency'], int):
157                surface.SDL_SetColorKey(imgsurface, True, image.info['transparency'])
158
159    return imgsurface
160