@@ -34,6 +34,27 @@ extern "C" {
3434#undef jmpbuf
3535#endif
3636
37+ struct buffer_t {
38+ PyObject *str;
39+ size_t cursor;
40+ size_t size;
41+ };
42+
43+
44+ static void write_png_data_buffer (png_structp png_ptr, png_bytep data, png_size_t length)
45+ {
46+ buffer_t *buff = (buffer_t *)png_get_io_ptr (png_ptr);
47+ if (buff->cursor + length < buff->size ) {
48+ memcpy (PyBytes_AS_STRING (buff->str ) + buff->cursor , data, length);
49+ buff->cursor += length;
50+ }
51+ }
52+
53+ static void flush_png_data_buffer (png_structp png_ptr)
54+ {
55+
56+ }
57+
3758static void write_png_data (png_structp png_ptr, png_bytep data, png_size_t length)
3859{
3960 PyObject *py_file_obj = (PyObject *)png_get_io_ptr (png_ptr);
@@ -69,20 +90,24 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
6990 numpy::array_view<unsigned char , 3 > buffer;
7091 PyObject *filein;
7192 double dpi = 0 ;
72- const char *names[] = { " buffer" , " file" , " dpi" , NULL };
93+ int compression = 6 ;
94+ int filter = -1 ;
95+ const char *names[] = { " buffer" , " file" , " dpi" , " compression" , " filter" , NULL };
7396
7497 // We don't need strict contiguity, just for each row to be
7598 // contiguous, and libpng has special handling for getting RGB out
7699 // of RGBA, ARGB or BGR. But the simplest thing to do is to
77100 // enforce contiguity using array_view::converter_contiguous.
78101 if (!PyArg_ParseTupleAndKeywords (args,
79102 kwds,
80- " O&O|d :write_png" ,
103+ " O&O|dii :write_png" ,
81104 (char **)names,
82105 &buffer.converter_contiguous ,
83106 &buffer,
84107 &filein,
85- &dpi)) {
108+ &dpi,
109+ &compression,
110+ &filter)) {
86111 return NULL ;
87112 }
88113
@@ -104,6 +129,8 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
104129 png_infop info_ptr = NULL ;
105130 struct png_color_8_struct sig_bit;
106131 int png_color_type;
132+ buffer_t buff;
133+ buff.str = NULL ;
107134
108135 switch (channels) {
109136 case 1 :
@@ -122,6 +149,12 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
122149 goto exit;
123150 }
124151
152+ if (compression < 0 || compression > 9 ) {
153+ PyErr_Format (PyExc_ValueError,
154+ " compression must be in range 0-9, got %d" , compression);
155+ goto exit;
156+ }
157+
125158 if (PyBytes_Check (filein) || PyUnicode_Check (filein)) {
126159 if ((py_file = mpl_PyFile_OpenFile (filein, (char *)" wb" )) == NULL ) {
127160 goto exit;
@@ -131,7 +164,14 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
131164 py_file = filein;
132165 }
133166
134- if ((fp = mpl_PyFile_Dup (py_file, (char *)" wb" , &offset))) {
167+ if (filein == Py_None) {
168+ buff.size = width * height * 4 + 1024 ;
169+ buff.str = PyBytes_FromStringAndSize (NULL , buff.size );
170+ if (buff.str == NULL ) {
171+ goto exit;
172+ }
173+ buff.cursor = 0 ;
174+ } else if ((fp = mpl_PyFile_Dup (py_file, (char *)" wb" , &offset))) {
135175 close_dup_file = true ;
136176 } else {
137177 PyErr_Clear ();
@@ -152,6 +192,11 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
152192 goto exit;
153193 }
154194
195+ png_set_compression_level (png_ptr, compression);
196+ if (filter >= 0 ) {
197+ png_set_filter (png_ptr, 0 , filter);
198+ }
199+
155200 info_ptr = png_create_info_struct (png_ptr);
156201 if (info_ptr == NULL ) {
157202 PyErr_SetString (PyExc_RuntimeError, " Could not create info struct" );
@@ -163,7 +208,9 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
163208 goto exit;
164209 }
165210
166- if (fp) {
211+ if (buff.str ) {
212+ png_set_write_fn (png_ptr, (void *)&buff, &write_png_data_buffer, &flush_png_data_buffer);
213+ } else if (fp) {
167214 png_init_io (png_ptr, fp);
168215 } else {
169216 png_set_write_fn (png_ptr, (void *)py_file, &write_png_data, &flush_png_data);
@@ -227,8 +274,13 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
227274 }
228275
229276 if (PyErr_Occurred ()) {
277+ Py_XDECREF (buff.str );
230278 return NULL ;
231279 } else {
280+ if (buff.str ) {
281+ _PyBytes_Resize (&buff.str , buff.cursor );
282+ return buff.str ;
283+ }
232284 Py_RETURN_NONE;
233285 }
234286}
@@ -558,6 +610,15 @@ extern "C" {
558610
559611 import_array ();
560612
613+ if (PyModule_AddIntConstant (m, " PNG_FILTER_NONE" , PNG_FILTER_NONE ) ||
614+ PyModule_AddIntConstant (m, " PNG_FILTER_SUB" , PNG_FILTER_SUB ) ||
615+ PyModule_AddIntConstant (m, " PNG_FILTER_UP" , PNG_FILTER_UP ) ||
616+ PyModule_AddIntConstant (m, " PNG_FILTER_AVG" , PNG_FILTER_AVG ) ||
617+ PyModule_AddIntConstant (m, " PNG_FILTER_PAETH" , PNG_FILTER_PAETH )) {
618+ INITERROR ;
619+ }
620+
621+
561622#if PY3K
562623 return m;
563624#endif
0 commit comments