You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

439 lines
14 KiB

4 years ago
  1. from qrcode import constants, exceptions, util
  2. from qrcode.image.base import BaseImage
  3. import six
  4. from bisect import bisect_left
  5. def make(data=None, **kwargs):
  6. qr = QRCode(**kwargs)
  7. qr.add_data(data)
  8. return qr.make_image()
  9. def _check_version(version):
  10. if version < 1 or version > 40:
  11. raise ValueError(
  12. "Invalid version (was %s, expected 1 to 40)" % version)
  13. def _check_box_size(size):
  14. if int(size) <= 0:
  15. raise ValueError(
  16. "Invalid box size (was %s, expected larger than 0)" % size)
  17. def _check_mask_pattern(mask_pattern):
  18. if mask_pattern is None:
  19. return
  20. if not isinstance(mask_pattern, int):
  21. raise TypeError(
  22. "Invalid mask pattern (was %s, expected int)" % type(mask_pattern))
  23. if mask_pattern < 0 or mask_pattern > 7:
  24. raise ValueError(
  25. "Mask pattern should be in range(8) (got %s)" % mask_pattern)
  26. class QRCode:
  27. def __init__(self, version=None,
  28. error_correction=constants.ERROR_CORRECT_M,
  29. box_size=10, border=4,
  30. image_factory=None,
  31. mask_pattern=None):
  32. _check_box_size(box_size)
  33. self.version = version and int(version)
  34. self.error_correction = int(error_correction)
  35. self.box_size = int(box_size)
  36. # Spec says border should be at least four boxes wide, but allow for
  37. # any (e.g. for producing printable QR codes).
  38. self.border = int(border)
  39. _check_mask_pattern(mask_pattern)
  40. self.mask_pattern = mask_pattern
  41. self.image_factory = image_factory
  42. if image_factory is not None:
  43. assert issubclass(image_factory, BaseImage)
  44. self.clear()
  45. def clear(self):
  46. """
  47. Reset the internal data.
  48. """
  49. self.modules = None
  50. self.modules_count = 0
  51. self.data_cache = None
  52. self.data_list = []
  53. def add_data(self, data, optimize=20):
  54. """
  55. Add data to this QR Code.
  56. :param optimize: Data will be split into multiple chunks to optimize
  57. the QR size by finding to more compressed modes of at least this
  58. length. Set to ``0`` to avoid optimizing at all.
  59. """
  60. if isinstance(data, util.QRData):
  61. self.data_list.append(data)
  62. else:
  63. if optimize:
  64. self.data_list.extend(
  65. util.optimal_data_chunks(data, minimum=optimize))
  66. else:
  67. self.data_list.append(util.QRData(data))
  68. self.data_cache = None
  69. def make(self, fit=True):
  70. """
  71. Compile the data into a QR Code array.
  72. :param fit: If ``True`` (or if a size has not been provided), find the
  73. best fit for the data to avoid data overflow errors.
  74. """
  75. if fit or (self.version is None):
  76. self.best_fit(start=self.version)
  77. if self.mask_pattern is None:
  78. self.makeImpl(False, self.best_mask_pattern())
  79. else:
  80. self.makeImpl(False, self.mask_pattern)
  81. def makeImpl(self, test, mask_pattern):
  82. _check_version(self.version)
  83. self.modules_count = self.version * 4 + 17
  84. self.modules = [None] * self.modules_count
  85. for row in range(self.modules_count):
  86. self.modules[row] = [None] * self.modules_count
  87. for col in range(self.modules_count):
  88. self.modules[row][col] = None # (col + row) % 3
  89. self.setup_position_probe_pattern(0, 0)
  90. self.setup_position_probe_pattern(self.modules_count - 7, 0)
  91. self.setup_position_probe_pattern(0, self.modules_count - 7)
  92. self.setup_position_adjust_pattern()
  93. self.setup_timing_pattern()
  94. self.setup_type_info(test, mask_pattern)
  95. if self.version >= 7:
  96. self.setup_type_number(test)
  97. if self.data_cache is None:
  98. self.data_cache = util.create_data(
  99. self.version, self.error_correction, self.data_list)
  100. self.map_data(self.data_cache, mask_pattern)
  101. def setup_position_probe_pattern(self, row, col):
  102. for r in range(-1, 8):
  103. if row + r <= -1 or self.modules_count <= row + r:
  104. continue
  105. for c in range(-1, 8):
  106. if col + c <= -1 or self.modules_count <= col + c:
  107. continue
  108. if (0 <= r and r <= 6 and (c == 0 or c == 6)
  109. or (0 <= c and c <= 6 and (r == 0 or r == 6))
  110. or (2 <= r and r <= 4 and 2 <= c and c <= 4)):
  111. self.modules[row + r][col + c] = True
  112. else:
  113. self.modules[row + r][col + c] = False
  114. def best_fit(self, start=None):
  115. """
  116. Find the minimum size required to fit in the data.
  117. """
  118. if start is None:
  119. start = 1
  120. _check_version(start)
  121. # Corresponds to the code in util.create_data, except we don't yet know
  122. # version, so optimistically assume start and check later
  123. mode_sizes = util.mode_sizes_for_version(start)
  124. buffer = util.BitBuffer()
  125. for data in self.data_list:
  126. buffer.put(data.mode, 4)
  127. buffer.put(len(data), mode_sizes[data.mode])
  128. data.write(buffer)
  129. needed_bits = len(buffer)
  130. self.version = bisect_left(util.BIT_LIMIT_TABLE[self.error_correction],
  131. needed_bits, start)
  132. if self.version == 41:
  133. raise exceptions.DataOverflowError()
  134. # Now check whether we need more bits for the mode sizes, recursing if
  135. # our guess was too low
  136. if mode_sizes is not util.mode_sizes_for_version(self.version):
  137. self.best_fit(start=self.version)
  138. return self.version
  139. def best_mask_pattern(self):
  140. """
  141. Find the most efficient mask pattern.
  142. """
  143. min_lost_point = 0
  144. pattern = 0
  145. for i in range(8):
  146. self.makeImpl(True, i)
  147. lost_point = util.lost_point(self.modules)
  148. if i == 0 or min_lost_point > lost_point:
  149. min_lost_point = lost_point
  150. pattern = i
  151. return pattern
  152. def print_tty(self, out=None):
  153. """
  154. Output the QR Code only using TTY colors.
  155. If the data has not been compiled yet, make it first.
  156. """
  157. if out is None:
  158. import sys
  159. out = sys.stdout
  160. if not out.isatty():
  161. raise OSError("Not a tty")
  162. if self.data_cache is None:
  163. self.make()
  164. modcount = self.modules_count
  165. out.write("\x1b[1;47m" + (" " * (modcount * 2 + 4)) + "\x1b[0m\n")
  166. for r in range(modcount):
  167. out.write("\x1b[1;47m \x1b[40m")
  168. for c in range(modcount):
  169. if self.modules[r][c]:
  170. out.write(" ")
  171. else:
  172. out.write("\x1b[1;47m \x1b[40m")
  173. out.write("\x1b[1;47m \x1b[0m\n")
  174. out.write("\x1b[1;47m" + (" " * (modcount * 2 + 4)) + "\x1b[0m\n")
  175. out.flush()
  176. def print_ascii(self, out=None, tty=False, invert=False):
  177. """
  178. Output the QR Code using ASCII characters.
  179. :param tty: use fixed TTY color codes (forces invert=True)
  180. :param invert: invert the ASCII characters (solid <-> transparent)
  181. """
  182. if out is None:
  183. import sys
  184. if sys.version_info < (2, 7):
  185. # On Python versions 2.6 and earlier, stdout tries to encode
  186. # strings using ASCII rather than stdout.encoding, so use this
  187. # workaround.
  188. import codecs
  189. out = codecs.getwriter(sys.stdout.encoding)(sys.stdout)
  190. else:
  191. out = sys.stdout
  192. if tty and not out.isatty():
  193. raise OSError("Not a tty")
  194. if self.data_cache is None:
  195. self.make()
  196. modcount = self.modules_count
  197. codes = [six.int2byte(code).decode('cp437')
  198. for code in (255, 223, 220, 219)]
  199. if tty:
  200. invert = True
  201. if invert:
  202. codes.reverse()
  203. def get_module(x, y):
  204. if (invert and self.border and
  205. max(x, y) >= modcount+self.border):
  206. return 1
  207. if min(x, y) < 0 or max(x, y) >= modcount:
  208. return 0
  209. return self.modules[x][y]
  210. for r in range(-self.border, modcount+self.border, 2):
  211. if tty:
  212. if not invert or r < modcount+self.border-1:
  213. out.write('\x1b[48;5;232m') # Background black
  214. out.write('\x1b[38;5;255m') # Foreground white
  215. for c in range(-self.border, modcount+self.border):
  216. pos = get_module(r, c) + (get_module(r+1, c) << 1)
  217. out.write(codes[pos])
  218. if tty:
  219. out.write('\x1b[0m')
  220. out.write('\n')
  221. out.flush()
  222. def make_image(self, image_factory=None, **kwargs):
  223. """
  224. Make an image from the QR Code data.
  225. If the data has not been compiled yet, make it first.
  226. """
  227. _check_box_size(self.box_size)
  228. if self.data_cache is None:
  229. self.make()
  230. if image_factory is not None:
  231. assert issubclass(image_factory, BaseImage)
  232. else:
  233. image_factory = self.image_factory
  234. if image_factory is None:
  235. # Use PIL by default
  236. from qrcode.image.pil import PilImage
  237. image_factory = PilImage
  238. im = image_factory(
  239. self.border, self.modules_count, self.box_size, **kwargs)
  240. for r in range(self.modules_count):
  241. for c in range(self.modules_count):
  242. if self.modules[r][c]:
  243. im.drawrect(r, c)
  244. return im
  245. def setup_timing_pattern(self):
  246. for r in range(8, self.modules_count - 8):
  247. if self.modules[r][6] is not None:
  248. continue
  249. self.modules[r][6] = (r % 2 == 0)
  250. for c in range(8, self.modules_count - 8):
  251. if self.modules[6][c] is not None:
  252. continue
  253. self.modules[6][c] = (c % 2 == 0)
  254. def setup_position_adjust_pattern(self):
  255. pos = util.pattern_position(self.version)
  256. for i in range(len(pos)):
  257. for j in range(len(pos)):
  258. row = pos[i]
  259. col = pos[j]
  260. if self.modules[row][col] is not None:
  261. continue
  262. for r in range(-2, 3):
  263. for c in range(-2, 3):
  264. if (r == -2 or r == 2 or c == -2 or c == 2 or
  265. (r == 0 and c == 0)):
  266. self.modules[row + r][col + c] = True
  267. else:
  268. self.modules[row + r][col + c] = False
  269. def setup_type_number(self, test):
  270. bits = util.BCH_type_number(self.version)
  271. for i in range(18):
  272. mod = (not test and ((bits >> i) & 1) == 1)
  273. self.modules[i // 3][i % 3 + self.modules_count - 8 - 3] = mod
  274. for i in range(18):
  275. mod = (not test and ((bits >> i) & 1) == 1)
  276. self.modules[i % 3 + self.modules_count - 8 - 3][i // 3] = mod
  277. def setup_type_info(self, test, mask_pattern):
  278. data = (self.error_correction << 3) | mask_pattern
  279. bits = util.BCH_type_info(data)
  280. # vertical
  281. for i in range(15):
  282. mod = (not test and ((bits >> i) & 1) == 1)
  283. if i < 6:
  284. self.modules[i][8] = mod
  285. elif i < 8:
  286. self.modules[i + 1][8] = mod
  287. else:
  288. self.modules[self.modules_count - 15 + i][8] = mod
  289. # horizontal
  290. for i in range(15):
  291. mod = (not test and ((bits >> i) & 1) == 1)
  292. if i < 8:
  293. self.modules[8][self.modules_count - i - 1] = mod
  294. elif i < 9:
  295. self.modules[8][15 - i - 1 + 1] = mod
  296. else:
  297. self.modules[8][15 - i - 1] = mod
  298. # fixed module
  299. self.modules[self.modules_count - 8][8] = (not test)
  300. def map_data(self, data, mask_pattern):
  301. inc = -1
  302. row = self.modules_count - 1
  303. bitIndex = 7
  304. byteIndex = 0
  305. mask_func = util.mask_func(mask_pattern)
  306. data_len = len(data)
  307. for col in six.moves.xrange(self.modules_count - 1, 0, -2):
  308. if col <= 6:
  309. col -= 1
  310. col_range = (col, col-1)
  311. while True:
  312. for c in col_range:
  313. if self.modules[row][c] is None:
  314. dark = False
  315. if byteIndex < data_len:
  316. dark = (((data[byteIndex] >> bitIndex) & 1) == 1)
  317. if mask_func(row, c):
  318. dark = not dark
  319. self.modules[row][c] = dark
  320. bitIndex -= 1
  321. if bitIndex == -1:
  322. byteIndex += 1
  323. bitIndex = 7
  324. row += inc
  325. if row < 0 or self.modules_count <= row:
  326. row -= inc
  327. inc = -inc
  328. break
  329. def get_matrix(self):
  330. """
  331. Return the QR Code as a multidimensonal array, including the border.
  332. To return the array without a border, set ``self.border`` to 0 first.
  333. """
  334. if self.data_cache is None:
  335. self.make()
  336. if not self.border:
  337. return self.modules
  338. width = len(self.modules) + self.border*2
  339. code = [[False]*width] * self.border
  340. x_border = [False]*self.border
  341. for module in self.modules:
  342. code.append(x_border + module + x_border)
  343. code += [[False]*width] * self.border
  344. return code