1# Copyright (c) 2013-2016 Jeffrey Pfau
 2#
 3# This Source Code Form is subject to the terms of the Mozilla Public
 4# License, v. 2.0. If a copy of the MPL was not distributed with this
 5# file, You can obtain one at http://mozilla.org/MPL/2.0/.
 6from ._pylib import ffi, lib  # pylint: disable=no-name-in-module
 7from . import vfs
 8
 9MODE_RGB = 0
10MODE_RGBA = 1
11MODE_INDEX = 2
12
13
14class PNG:
15    def __init__(self, f, mode=MODE_RGB):
16        self._vfile = vfs.open(f)
17        self._png = None
18        self._info = None
19        self.mode = mode
20
21    def write_header(self, image):
22        self._png = lib.PNGWriteOpen(self._vfile.handle)
23        if self.mode == MODE_RGB:
24            self._info = lib.PNGWriteHeader(self._png, image.width, image.height)
25        if self.mode == MODE_RGBA:
26            self._info = lib.PNGWriteHeaderA(self._png, image.width, image.height)
27        if self.mode == MODE_INDEX:
28            self._info = lib.PNGWriteHeader8(self._png, image.width, image.height)
29        return self._info != ffi.NULL
30
31    def write_pixels(self, image):
32        if self.mode == MODE_RGB:
33            return lib.PNGWritePixels(self._png, image.width, image.height, image.stride, image.buffer)
34        if self.mode == MODE_RGBA:
35            return lib.PNGWritePixelsA(self._png, image.width, image.height, image.stride, image.buffer)
36        if self.mode == MODE_INDEX:
37            return lib.PNGWritePixels8(self._png, image.width, image.height, image.stride, image.buffer)
38        return False
39
40    def write_close(self):
41        lib.PNGWriteClose(self._png, self._info)
42        self._png = None
43        self._info = None