1import io
2import os
3import sys
4import unittest
5import warnings
6from tempfile import mkdtemp
7from unittest import mock
8
9try:
10    import pymaging_png  # ensure that PNG support is installed
11    import qrcode.image.pure
12except ImportError:  # pragma: no cover
13    pymaging_png = None
14
15import qrcode
16import qrcode.image.svg
17import qrcode.util
18from qrcode.exceptions import DataOverflowError
19from qrcode.image.base import BaseImage
20from qrcode.image.pil import Image as pil_Image
21from qrcode.image.styledpil import StyledPilImage
22from qrcode.image.styles import colormasks, moduledrawers
23from qrcode.tests.svg import SvgImageWhite
24from qrcode.util import MODE_8BIT_BYTE, MODE_ALPHA_NUM, MODE_NUMBER, QRData
25
26UNICODE_TEXT = '\u03b1\u03b2\u03b3'
27WHITE = (255, 255, 255)
28BLACK = (0, 0, 0)
29RED = (255, 0, 0)
30
31
32class QRCodeTests(unittest.TestCase):
33    def setUp(self):
34        self.tmpdir = mkdtemp()
35
36    def tearDown(self):
37        os.rmdir(self.tmpdir)
38
39    def test_basic(self):
40        qr = qrcode.QRCode(version=1)
41        qr.add_data('a')
42        qr.make(fit=False)
43
44    def test_large(self):
45        qr = qrcode.QRCode(version=27)
46        qr.add_data('a')
47        qr.make(fit=False)
48
49    def test_invalid_version(self):
50        qr = qrcode.QRCode(version=41)
51        self.assertRaises(ValueError, qr.make, fit=False)
52
53    def test_invalid_border(self):
54        self.assertRaises(ValueError, qrcode.QRCode, border=-1)
55
56    def test_overflow(self):
57        qr = qrcode.QRCode(version=1)
58        qr.add_data('abcdefghijklmno')
59        self.assertRaises(DataOverflowError, qr.make, fit=False)
60
61    def test_add_qrdata(self):
62        qr = qrcode.QRCode(version=1)
63        data = QRData('a')
64        qr.add_data(data)
65        qr.make(fit=False)
66
67    def test_fit(self):
68        qr = qrcode.QRCode()
69        qr.add_data('a')
70        qr.make()
71        self.assertEqual(qr.version, 1)
72        qr.add_data('bcdefghijklmno')
73        qr.make()
74        self.assertEqual(qr.version, 2)
75
76    def test_mode_number(self):
77        qr = qrcode.QRCode()
78        qr.add_data('1234567890123456789012345678901234', optimize=0)
79        qr.make()
80        self.assertEqual(qr.version, 1)
81        self.assertEqual(qr.data_list[0].mode, MODE_NUMBER)
82
83    def test_mode_alpha(self):
84        qr = qrcode.QRCode()
85        qr.add_data('ABCDEFGHIJ1234567890', optimize=0)
86        qr.make()
87        self.assertEqual(qr.version, 1)
88        self.assertEqual(qr.data_list[0].mode, MODE_ALPHA_NUM)
89
90    def test_regression_mode_comma(self):
91        qr = qrcode.QRCode()
92        qr.add_data(',', optimize=0)
93        qr.make()
94        self.assertEqual(qr.data_list[0].mode, MODE_8BIT_BYTE)
95
96    def test_mode_8bit(self):
97        qr = qrcode.QRCode()
98        qr.add_data('abcABC' + UNICODE_TEXT, optimize=0)
99        qr.make()
100        self.assertEqual(qr.version, 1)
101        self.assertEqual(qr.data_list[0].mode, MODE_8BIT_BYTE)
102
103    def test_mode_8bit_newline(self):
104        qr = qrcode.QRCode()
105        qr.add_data('ABCDEFGHIJ1234567890\n', optimize=0)
106        qr.make()
107        self.assertEqual(qr.data_list[0].mode, MODE_8BIT_BYTE)
108
109    def test_render_pil(self):
110        qr = qrcode.QRCode()
111        qr.add_data(UNICODE_TEXT)
112        img = qr.make_image()
113        img.save(io.BytesIO())
114        self.assertIsInstance(img.get_image(), pil_Image.Image)
115
116    def test_render_pil_with_transparent_background(self):
117        qr = qrcode.QRCode()
118        qr.add_data(UNICODE_TEXT)
119        img = qr.make_image(back_color='TransParent')
120        img.save(io.BytesIO())
121
122    def test_render_pil_with_red_background(self):
123        qr = qrcode.QRCode()
124        qr.add_data(UNICODE_TEXT)
125        img = qr.make_image(back_color='red')
126        img.save(io.BytesIO())
127
128    def test_render_pil_with_rgb_color_tuples(self):
129        qr = qrcode.QRCode()
130        qr.add_data(UNICODE_TEXT)
131        img = qr.make_image(back_color=(255, 195, 235), fill_color=(55, 95, 35))
132        img.save(io.BytesIO())
133
134    def test_render_with_pattern(self):
135        qr = qrcode.QRCode(mask_pattern=3)
136        qr.add_data(UNICODE_TEXT)
137        img = qr.make_image()
138        img.save(io.BytesIO())
139
140    def test_make_image_with_wrong_pattern(self):
141        with self.assertRaises(TypeError):
142            qrcode.QRCode(mask_pattern='string pattern')
143
144        with self.assertRaises(ValueError):
145            qrcode.QRCode(mask_pattern=-1)
146
147        with self.assertRaises(ValueError):
148            qrcode.QRCode(mask_pattern=42)
149
150    def test_mask_pattern_setter(self):
151        qr = qrcode.QRCode()
152
153        with self.assertRaises(TypeError):
154            qr.mask_pattern = "string pattern"
155
156        with self.assertRaises(ValueError):
157            qr.mask_pattern = -1
158
159        with self.assertRaises(ValueError):
160            qr.mask_pattern = 8
161
162    def test_qrcode_bad_factory(self):
163        with self.assertRaises(TypeError):
164           qrcode.QRCode(image_factory='not_BaseImage')
165
166        with self.assertRaises(AssertionError):
167            qrcode.QRCode(image_factory=dict)
168
169    def test_qrcode_factory(self):
170
171        class MockFactory(BaseImage):
172            drawrect = mock.Mock()
173            new_image = mock.Mock()
174
175        qr = qrcode.QRCode(image_factory=MockFactory)
176        qr.add_data(UNICODE_TEXT)
177        qr.make_image()
178        self.assertTrue(MockFactory.new_image.called)
179        self.assertTrue(MockFactory.drawrect.called)
180
181    def test_render_svg(self):
182        qr = qrcode.QRCode()
183        qr.add_data(UNICODE_TEXT)
184        img = qr.make_image(image_factory=qrcode.image.svg.SvgImage)
185        img.save(io.BytesIO())
186
187    def test_render_svg_path(self):
188        qr = qrcode.QRCode()
189        qr.add_data(UNICODE_TEXT)
190        img = qr.make_image(image_factory=qrcode.image.svg.SvgPathImage)
191        img.save(io.BytesIO())
192
193    def test_render_svg_fragment(self):
194        qr = qrcode.QRCode()
195        qr.add_data(UNICODE_TEXT)
196        img = qr.make_image(image_factory=qrcode.image.svg.SvgFragmentImage)
197        img.save(io.BytesIO())
198
199    def test_svg_string(self):
200        qr = qrcode.QRCode()
201        qr.add_data(UNICODE_TEXT)
202        img = qr.make_image(image_factory=qrcode.image.svg.SvgFragmentImage)
203        file_like = io.BytesIO()
204        img.save(file_like)
205        file_like.seek(0)
206        assert file_like.read() in img.to_string()
207
208    def test_render_svg_with_background(self):
209        qr = qrcode.QRCode()
210        qr.add_data(UNICODE_TEXT)
211        img = qr.make_image(image_factory=SvgImageWhite)
212        img.save(io.BytesIO())
213
214    @unittest.skipIf(not pymaging_png, "Requires pymaging with PNG support")
215    def test_render_pymaging_png(self):
216        qr = qrcode.QRCode()
217        qr.add_data(UNICODE_TEXT)
218        img = qr.make_image(image_factory=qrcode.image.pure.PymagingImage)
219        from pymaging import Image as pymaging_Image
220        self.assertIsInstance(img.get_image(), pymaging_Image)
221        with warnings.catch_warnings():
222            warnings.simplefilter('ignore', DeprecationWarning)
223            img.save(io.BytesIO())
224
225    @unittest.skipIf(not pymaging_png, "Requires pymaging")
226    def test_render_pymaging_png_bad_kind(self):
227        qr = qrcode.QRCode()
228        qr.add_data(UNICODE_TEXT)
229        img = qr.make_image(image_factory=qrcode.image.pure.PymagingImage)
230        with self.assertRaises(ValueError):
231            img.save(io.BytesIO(), kind='FISH')
232
233    def test_render_styled_pil_image(self):
234        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
235        qr.add_data(UNICODE_TEXT)
236        img = qr.make_image(image_factory=StyledPilImage)
237        img.save(io.BytesIO())
238
239    def test_render_styled_with_embeded_image(self):
240        embeded_img = pil_Image.new('RGB', (10, 10), color='red')
241        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
242        qr.add_data(UNICODE_TEXT)
243        img = qr.make_image(image_factory=StyledPilImage, embeded_image=embeded_img)
244        img.save(io.BytesIO())
245
246    def test_render_styled_with_embeded_image_path(self):
247        tmpfile = os.path.join(self.tmpdir, "test.png")
248        embeded_img = pil_Image.new('RGB', (10, 10), color='red')
249        embeded_img.save(tmpfile)
250        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
251        qr.add_data(UNICODE_TEXT)
252        img = qr.make_image(image_factory=StyledPilImage, embeded_image_path=tmpfile)
253        img.save(io.BytesIO())
254        os.remove(tmpfile)
255
256    def test_render_styled_with_square_module_drawer(self):
257        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
258        qr.add_data(UNICODE_TEXT)
259        img = qr.make_image(image_factory=StyledPilImage, module_drawer=moduledrawers.SquareModuleDrawer())
260        img.save(io.BytesIO())
261
262    def test_render_styled_with_gapped_module_drawer(self):
263        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
264        qr.add_data(UNICODE_TEXT)
265        img = qr.make_image(image_factory=StyledPilImage, module_drawer=moduledrawers.GappedSquareModuleDrawer())
266        img.save(io.BytesIO())
267
268    def test_render_styled_with_circle_module_drawer(self):
269        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
270        qr.add_data(UNICODE_TEXT)
271        img = qr.make_image(image_factory=StyledPilImage, module_drawer=moduledrawers.CircleModuleDrawer())
272        img.save(io.BytesIO())
273
274    def test_render_styled_with_rounded_module_drawer(self):
275        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
276        qr.add_data(UNICODE_TEXT)
277        img = qr.make_image(image_factory=StyledPilImage, module_drawer=moduledrawers.RoundedModuleDrawer())
278        img.save(io.BytesIO())
279
280    def test_render_styled_with_vertical_bars_module_drawer(self):
281        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
282        qr.add_data(UNICODE_TEXT)
283        img = qr.make_image(image_factory=StyledPilImage, module_drawer=moduledrawers.VerticalBarsDrawer())
284        img.save(io.BytesIO())
285
286    def test_render_styled_with_horizontal_bars_module_drawer(self):
287        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
288        qr.add_data(UNICODE_TEXT)
289        img = qr.make_image(image_factory=StyledPilImage, module_drawer=moduledrawers.HorizontalBarsDrawer())
290        img.save(io.BytesIO())
291
292    def test_render_styled_with_default_solid_color_mask(self):
293        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
294        qr.add_data(UNICODE_TEXT)
295        mask = colormasks.SolidFillColorMask()
296        img = qr.make_image(image_factory=StyledPilImage, color_mask=mask)
297        img.save(io.BytesIO())
298
299    def test_render_styled_with_solid_color_mask(self):
300        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
301        qr.add_data(UNICODE_TEXT)
302        mask = colormasks.SolidFillColorMask(back_color=WHITE, front_color=RED)
303        img = qr.make_image(image_factory=StyledPilImage, color_mask=mask)
304        img.save(io.BytesIO())
305
306    def test_render_styled_with_color_mask_with_transparency(self):
307        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
308        qr.add_data(UNICODE_TEXT)
309        mask = colormasks.SolidFillColorMask(back_color=(255, 0, 255, 255), front_color=RED)
310        img = qr.make_image(image_factory=StyledPilImage, color_mask=mask)
311        img.save(io.BytesIO())
312        assert img.mode == "RGBA"
313
314    def test_render_styled_with_radial_gradient_color_mask(self):
315        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
316        qr.add_data(UNICODE_TEXT)
317        mask = colormasks.RadialGradiantColorMask(back_color=WHITE, center_color=BLACK, edge_color=RED)
318        img = qr.make_image(image_factory=StyledPilImage, color_mask=mask)
319        img.save(io.BytesIO())
320
321    def test_render_styled_with_square_gradient_color_mask(self):
322        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
323        qr.add_data(UNICODE_TEXT)
324        mask = colormasks.SquareGradiantColorMask(back_color=WHITE, center_color=BLACK, edge_color=RED)
325        img = qr.make_image(image_factory=StyledPilImage, color_mask=mask)
326        img.save(io.BytesIO())
327
328    def test_render_styled_with_horizontal_gradient_color_mask(self):
329        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
330        qr.add_data(UNICODE_TEXT)
331        mask = colormasks.HorizontalGradiantColorMask(back_color=WHITE, left_color=RED, right_color=BLACK)
332        img = qr.make_image(image_factory=StyledPilImage, color_mask=mask)
333        img.save(io.BytesIO())
334
335    def test_render_styled_with_vertical_gradient_color_mask(self):
336        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
337        qr.add_data(UNICODE_TEXT)
338        mask = colormasks.VerticalGradiantColorMask(back_color=WHITE, top_color=RED, bottom_color=BLACK)
339        img = qr.make_image(image_factory=StyledPilImage, color_mask=mask)
340        img.save(io.BytesIO())
341
342    def test_render_styled_with_image_color_mask(self):
343        img_mask = pil_Image.new('RGB', (10, 10), color='red')
344        qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_L)
345        qr.add_data(UNICODE_TEXT)
346        mask = colormasks.ImageColorMask(back_color=WHITE, color_mask_image=img_mask)
347        img = qr.make_image(image_factory=StyledPilImage, color_mask=mask)
348        img.save(io.BytesIO())
349
350    def test_optimize(self):
351        qr = qrcode.QRCode()
352        text = 'A1abc12345def1HELLOa'
353        qr.add_data(text, optimize=4)
354        qr.make()
355        self.assertEqual(
356            [d.mode for d in qr.data_list],
357            [
358                MODE_8BIT_BYTE, MODE_NUMBER, MODE_8BIT_BYTE, MODE_ALPHA_NUM,
359                MODE_8BIT_BYTE
360            ]
361        )
362        self.assertEqual(qr.version, 2)
363
364    def test_optimize_short(self):
365        qr = qrcode.QRCode()
366        text = 'A1abc1234567def1HELLOa'
367        qr.add_data(text, optimize=7)
368        qr.make()
369        self.assertEqual(len(qr.data_list), 3)
370        self.assertEqual(
371            [d.mode for d in qr.data_list],
372            [MODE_8BIT_BYTE, MODE_NUMBER, MODE_8BIT_BYTE]
373        )
374        self.assertEqual(qr.version, 2)
375
376    def test_optimize_longer_than_data(self):
377        qr = qrcode.QRCode()
378        text = 'ABCDEFGHIJK'
379        qr.add_data(text, optimize=12)
380        self.assertEqual(len(qr.data_list), 1)
381        self.assertEqual(qr.data_list[0].mode, MODE_ALPHA_NUM)
382
383    def test_optimize_size(self):
384        text = 'A1abc12345123451234512345def1HELLOHELLOHELLOHELLOa' * 5
385
386        qr = qrcode.QRCode()
387        qr.add_data(text)
388        qr.make()
389        self.assertEqual(qr.version, 10)
390
391        qr = qrcode.QRCode()
392        qr.add_data(text, optimize=0)
393        qr.make()
394        self.assertEqual(qr.version, 11)
395
396    def test_qrdata_repr(self):
397        data = b'hello'
398        data_obj = qrcode.util.QRData(data)
399        self.assertEqual(repr(data_obj), repr(data))
400
401    def test_print_ascii_stdout(self):
402        qr = qrcode.QRCode()
403        stdout_encoding = sys.stdout.encoding
404        with mock.patch('sys.stdout') as fake_stdout:
405            # Python 2.6 needs sys.stdout.encoding to be a real string.
406            sys.stdout.encoding = stdout_encoding
407            fake_stdout.isatty.return_value = None
408            self.assertRaises(OSError, qr.print_ascii, tty=True)
409            self.assertTrue(fake_stdout.isatty.called)
410
411    def test_print_ascii(self):
412        qr = qrcode.QRCode(border=0)
413        f = io.StringIO()
414        qr.print_ascii(out=f)
415        printed = f.getvalue()
416        f.close()
417        expected = '\u2588\u2580\u2580\u2580\u2580\u2580\u2588'
418        self.assertEqual(printed[:len(expected)], expected)
419
420        f = io.StringIO()
421        f.isatty = lambda: True
422        qr.print_ascii(out=f, tty=True)
423        printed = f.getvalue()
424        f.close()
425        expected = (
426            '\x1b[48;5;232m\x1b[38;5;255m' +
427            '\xa0\u2584\u2584\u2584\u2584\u2584\xa0')
428        self.assertEqual(printed[:len(expected)], expected)
429
430    def test_print_tty_stdout(self):
431        qr = qrcode.QRCode()
432        with mock.patch('sys.stdout') as fake_stdout:
433            fake_stdout.isatty.return_value = None
434            self.assertRaises(OSError, qr.print_tty)
435            self.assertTrue(fake_stdout.isatty.called)
436
437    def test_print_tty(self):
438        qr = qrcode.QRCode()
439        f = io.StringIO()
440        f.isatty = lambda: True
441        qr.print_tty(out=f)
442        printed = f.getvalue()
443        f.close()
444        BOLD_WHITE_BG = '\x1b[1;47m'
445        BLACK_BG = '\x1b[40m'
446        WHITE_BLOCK = BOLD_WHITE_BG + '  ' + BLACK_BG
447        EOL = '\x1b[0m\n'
448        expected = (
449            BOLD_WHITE_BG + '  '*23 + EOL +
450            WHITE_BLOCK + '  '*7 + WHITE_BLOCK)
451        self.assertEqual(printed[:len(expected)], expected)
452
453    def test_get_matrix(self):
454        qr = qrcode.QRCode(border=0)
455        qr.add_data('1')
456        self.assertEqual(qr.get_matrix(), qr.modules)
457
458    def test_get_matrix_border(self):
459        qr = qrcode.QRCode(border=1)
460        qr.add_data('1')
461        matrix = [row[1:-1] for row in qr.get_matrix()[1:-1]]
462        self.assertEqual(matrix, qr.modules)
463
464    def test_negative_size_at_construction(self):
465        self.assertRaises(ValueError, qrcode.QRCode, box_size=-1)
466
467    def test_negative_size_at_usage(self):
468        qr = qrcode.QRCode()
469        qr.box_size = -1
470        self.assertRaises(ValueError, qr.make_image)
471
472
473class ShortcutTest(unittest.TestCase):
474
475    def runTest(self):
476        qrcode.make('image')
477