main.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541
  1. import sys
  2. from bisect import bisect_left
  3. from typing import (
  4. Generic,
  5. NamedTuple,
  6. Optional,
  7. TypeVar,
  8. cast,
  9. overload,
  10. Literal,
  11. )
  12. from qrcode import constants, exceptions, util
  13. from qrcode.image.base import BaseImage
  14. from qrcode.image.pure import PyPNGImage
  15. ModulesType = list[list[Optional[bool]]]
  16. # Cache modules generated just based on the QR Code version
  17. precomputed_qr_blanks: dict[int, ModulesType] = {}
  18. def make(data=None, **kwargs):
  19. qr = QRCode(**kwargs)
  20. qr.add_data(data)
  21. return qr.make_image()
  22. def _check_box_size(size):
  23. if int(size) <= 0:
  24. raise ValueError(f"Invalid box size (was {size}, expected larger than 0)")
  25. def _check_border(size):
  26. if int(size) < 0:
  27. raise ValueError(
  28. "Invalid border value (was %s, expected 0 or larger than that)" % size
  29. )
  30. def _check_mask_pattern(mask_pattern):
  31. if mask_pattern is None:
  32. return
  33. if not isinstance(mask_pattern, int):
  34. raise TypeError(
  35. f"Invalid mask pattern (was {type(mask_pattern)}, expected int)"
  36. )
  37. if mask_pattern < 0 or mask_pattern > 7:
  38. raise ValueError(f"Mask pattern should be in range(8) (got {mask_pattern})")
  39. def copy_2d_array(x):
  40. return [row[:] for row in x]
  41. class ActiveWithNeighbors(NamedTuple):
  42. NW: bool
  43. N: bool
  44. NE: bool
  45. W: bool
  46. me: bool
  47. E: bool
  48. SW: bool
  49. S: bool
  50. SE: bool
  51. def __bool__(self) -> bool:
  52. return self.me
  53. GenericImage = TypeVar("GenericImage", bound=BaseImage)
  54. GenericImageLocal = TypeVar("GenericImageLocal", bound=BaseImage)
  55. class QRCode(Generic[GenericImage]):
  56. modules: ModulesType
  57. _version: Optional[int] = None
  58. def __init__(
  59. self,
  60. version=None,
  61. error_correction=constants.ERROR_CORRECT_M,
  62. box_size=10,
  63. border=4,
  64. image_factory: Optional[type[GenericImage]] = None,
  65. mask_pattern=None,
  66. ):
  67. _check_box_size(box_size)
  68. _check_border(border)
  69. self.version = version
  70. self.error_correction = int(error_correction)
  71. self.box_size = int(box_size)
  72. # Spec says border should be at least four boxes wide, but allow for
  73. # any (e.g. for producing printable QR codes).
  74. self.border = int(border)
  75. self.mask_pattern = mask_pattern
  76. self.image_factory = image_factory
  77. if image_factory is not None:
  78. assert issubclass(image_factory, BaseImage)
  79. self.clear()
  80. @property
  81. def version(self) -> int:
  82. if self._version is None:
  83. self.best_fit()
  84. return cast(int, self._version)
  85. @version.setter
  86. def version(self, value) -> None:
  87. if value is not None:
  88. value = int(value)
  89. util.check_version(value)
  90. self._version = value
  91. @property
  92. def mask_pattern(self):
  93. return self._mask_pattern
  94. @mask_pattern.setter
  95. def mask_pattern(self, pattern):
  96. _check_mask_pattern(pattern)
  97. self._mask_pattern = pattern
  98. def clear(self):
  99. """
  100. Reset the internal data.
  101. """
  102. self.modules = [[]]
  103. self.modules_count = 0
  104. self.data_cache = None
  105. self.data_list = []
  106. def add_data(self, data, optimize=20):
  107. """
  108. Add data to this QR Code.
  109. :param optimize: Data will be split into multiple chunks to optimize
  110. the QR size by finding to more compressed modes of at least this
  111. length. Set to ``0`` to avoid optimizing at all.
  112. """
  113. if isinstance(data, util.QRData):
  114. self.data_list.append(data)
  115. elif optimize:
  116. self.data_list.extend(util.optimal_data_chunks(data, minimum=optimize))
  117. else:
  118. self.data_list.append(util.QRData(data))
  119. self.data_cache = None
  120. def make(self, fit=True):
  121. """
  122. Compile the data into a QR Code array.
  123. :param fit: If ``True`` (or if a size has not been provided), find the
  124. best fit for the data to avoid data overflow errors.
  125. """
  126. if fit or (self.version is None):
  127. self.best_fit(start=self.version)
  128. if self.mask_pattern is None:
  129. self.makeImpl(False, self.best_mask_pattern())
  130. else:
  131. self.makeImpl(False, self.mask_pattern)
  132. def makeImpl(self, test, mask_pattern):
  133. self.modules_count = self.version * 4 + 17
  134. if self.version in precomputed_qr_blanks:
  135. self.modules = copy_2d_array(precomputed_qr_blanks[self.version])
  136. else:
  137. self.modules = [
  138. [None] * self.modules_count for i in range(self.modules_count)
  139. ]
  140. self.setup_position_probe_pattern(0, 0)
  141. self.setup_position_probe_pattern(self.modules_count - 7, 0)
  142. self.setup_position_probe_pattern(0, self.modules_count - 7)
  143. self.setup_position_adjust_pattern()
  144. self.setup_timing_pattern()
  145. precomputed_qr_blanks[self.version] = copy_2d_array(self.modules)
  146. self.setup_type_info(test, mask_pattern)
  147. if self.version >= 7:
  148. self.setup_type_number(test)
  149. if self.data_cache is None:
  150. self.data_cache = util.create_data(
  151. self.version, self.error_correction, self.data_list
  152. )
  153. self.map_data(self.data_cache, mask_pattern)
  154. def setup_position_probe_pattern(self, row, col):
  155. for r in range(-1, 8):
  156. if row + r <= -1 or self.modules_count <= row + r:
  157. continue
  158. for c in range(-1, 8):
  159. if col + c <= -1 or self.modules_count <= col + c:
  160. continue
  161. if (
  162. (0 <= r <= 6 and c in {0, 6})
  163. or (0 <= c <= 6 and r in {0, 6})
  164. or (2 <= r <= 4 and 2 <= c <= 4)
  165. ):
  166. self.modules[row + r][col + c] = True
  167. else:
  168. self.modules[row + r][col + c] = False
  169. def best_fit(self, start=None):
  170. """
  171. Find the minimum size required to fit in the data.
  172. """
  173. if start is None:
  174. start = 1
  175. util.check_version(start)
  176. # Corresponds to the code in util.create_data, except we don't yet know
  177. # version, so optimistically assume start and check later
  178. mode_sizes = util.mode_sizes_for_version(start)
  179. buffer = util.BitBuffer()
  180. for data in self.data_list:
  181. buffer.put(data.mode, 4)
  182. buffer.put(len(data), mode_sizes[data.mode])
  183. data.write(buffer)
  184. needed_bits = len(buffer)
  185. self.version = bisect_left(
  186. util.BIT_LIMIT_TABLE[self.error_correction], needed_bits, start
  187. )
  188. if self.version == 41:
  189. raise exceptions.DataOverflowError()
  190. # Now check whether we need more bits for the mode sizes, recursing if
  191. # our guess was too low
  192. if mode_sizes is not util.mode_sizes_for_version(self.version):
  193. self.best_fit(start=self.version)
  194. return self.version
  195. def best_mask_pattern(self):
  196. """
  197. Find the most efficient mask pattern.
  198. """
  199. min_lost_point = 0
  200. pattern = 0
  201. for i in range(8):
  202. self.makeImpl(True, i)
  203. lost_point = util.lost_point(self.modules)
  204. if i == 0 or min_lost_point > lost_point:
  205. min_lost_point = lost_point
  206. pattern = i
  207. return pattern
  208. def print_tty(self, out=None):
  209. """
  210. Output the QR Code only using TTY colors.
  211. If the data has not been compiled yet, make it first.
  212. """
  213. if out is None:
  214. import sys
  215. out = sys.stdout
  216. if not out.isatty():
  217. raise OSError("Not a tty")
  218. if self.data_cache is None:
  219. self.make()
  220. modcount = self.modules_count
  221. out.write("\x1b[1;47m" + (" " * (modcount * 2 + 4)) + "\x1b[0m\n")
  222. for r in range(modcount):
  223. out.write("\x1b[1;47m \x1b[40m")
  224. for c in range(modcount):
  225. if self.modules[r][c]:
  226. out.write(" ")
  227. else:
  228. out.write("\x1b[1;47m \x1b[40m")
  229. out.write("\x1b[1;47m \x1b[0m\n")
  230. out.write("\x1b[1;47m" + (" " * (modcount * 2 + 4)) + "\x1b[0m\n")
  231. out.flush()
  232. def print_ascii(self, out=None, tty=False, invert=False):
  233. """
  234. Output the QR Code using ASCII characters.
  235. :param tty: use fixed TTY color codes (forces invert=True)
  236. :param invert: invert the ASCII characters (solid <-> transparent)
  237. """
  238. if out is None:
  239. out = sys.stdout
  240. if tty and not out.isatty():
  241. raise OSError("Not a tty")
  242. if self.data_cache is None:
  243. self.make()
  244. modcount = self.modules_count
  245. codes = [bytes((code,)).decode("cp437") for code in (255, 223, 220, 219)]
  246. if tty:
  247. invert = True
  248. if invert:
  249. codes.reverse()
  250. def get_module(x, y) -> int:
  251. if invert and self.border and max(x, y) >= modcount + self.border:
  252. return 1
  253. if min(x, y) < 0 or max(x, y) >= modcount:
  254. return 0
  255. return cast(int, self.modules[x][y])
  256. for r in range(-self.border, modcount + self.border, 2):
  257. if tty:
  258. if not invert or r < modcount + self.border - 1:
  259. out.write("\x1b[48;5;232m") # Background black
  260. out.write("\x1b[38;5;255m") # Foreground white
  261. for c in range(-self.border, modcount + self.border):
  262. pos = get_module(r, c) + (get_module(r + 1, c) << 1)
  263. out.write(codes[pos])
  264. if tty:
  265. out.write("\x1b[0m")
  266. out.write("\n")
  267. out.flush()
  268. @overload
  269. def make_image(
  270. self, image_factory: Literal[None] = None, **kwargs
  271. ) -> GenericImage: ...
  272. @overload
  273. def make_image(
  274. self, image_factory: type[GenericImageLocal] = None, **kwargs
  275. ) -> GenericImageLocal: ...
  276. def make_image(self, image_factory=None, **kwargs):
  277. """
  278. Make an image from the QR Code data.
  279. If the data has not been compiled yet, make it first.
  280. """
  281. # allow embeded_ parameters with typos for backwards compatibility
  282. if (
  283. kwargs.get("embedded_image_path")
  284. or kwargs.get("embedded_image")
  285. or kwargs.get("embeded_image_path")
  286. or kwargs.get("embeded_image")
  287. ) and self.error_correction != constants.ERROR_CORRECT_H:
  288. raise ValueError(
  289. "Error correction level must be ERROR_CORRECT_H if an embedded image is provided"
  290. )
  291. _check_box_size(self.box_size)
  292. if self.data_cache is None:
  293. self.make()
  294. if image_factory is not None:
  295. assert issubclass(image_factory, BaseImage)
  296. else:
  297. image_factory = self.image_factory
  298. if image_factory is None:
  299. from qrcode.image.pil import Image, PilImage
  300. # Use PIL by default if available, otherwise use PyPNG.
  301. image_factory = PilImage if Image else PyPNGImage
  302. im = image_factory(
  303. self.border,
  304. self.modules_count,
  305. self.box_size,
  306. qrcode_modules=self.modules,
  307. **kwargs,
  308. )
  309. if im.needs_drawrect:
  310. for r in range(self.modules_count):
  311. for c in range(self.modules_count):
  312. if im.needs_context:
  313. im.drawrect_context(r, c, qr=self)
  314. elif self.modules[r][c]:
  315. im.drawrect(r, c)
  316. if im.needs_processing:
  317. im.process()
  318. return im
  319. # return true if and only if (row, col) is in the module
  320. def is_constrained(self, row: int, col: int) -> bool:
  321. return (
  322. row >= 0
  323. and row < len(self.modules)
  324. and col >= 0
  325. and col < len(self.modules[row])
  326. )
  327. def setup_timing_pattern(self):
  328. for r in range(8, self.modules_count - 8):
  329. if self.modules[r][6] is not None:
  330. continue
  331. self.modules[r][6] = r % 2 == 0
  332. for c in range(8, self.modules_count - 8):
  333. if self.modules[6][c] is not None:
  334. continue
  335. self.modules[6][c] = c % 2 == 0
  336. def setup_position_adjust_pattern(self):
  337. pos = util.pattern_position(self.version)
  338. for i in range(len(pos)):
  339. row = pos[i]
  340. for j in range(len(pos)):
  341. col = pos[j]
  342. if self.modules[row][col] is not None:
  343. continue
  344. for r in range(-2, 3):
  345. for c in range(-2, 3):
  346. if (
  347. r == -2
  348. or r == 2
  349. or c == -2
  350. or c == 2
  351. or (r == 0 and c == 0)
  352. ):
  353. self.modules[row + r][col + c] = True
  354. else:
  355. self.modules[row + r][col + c] = False
  356. def setup_type_number(self, test):
  357. bits = util.BCH_type_number(self.version)
  358. for i in range(18):
  359. mod = not test and ((bits >> i) & 1) == 1
  360. self.modules[i // 3][i % 3 + self.modules_count - 8 - 3] = mod
  361. for i in range(18):
  362. mod = not test and ((bits >> i) & 1) == 1
  363. self.modules[i % 3 + self.modules_count - 8 - 3][i // 3] = mod
  364. def setup_type_info(self, test, mask_pattern):
  365. data = (self.error_correction << 3) | mask_pattern
  366. bits = util.BCH_type_info(data)
  367. # vertical
  368. for i in range(15):
  369. mod = not test and ((bits >> i) & 1) == 1
  370. if i < 6:
  371. self.modules[i][8] = mod
  372. elif i < 8:
  373. self.modules[i + 1][8] = mod
  374. else:
  375. self.modules[self.modules_count - 15 + i][8] = mod
  376. # horizontal
  377. for i in range(15):
  378. mod = not test and ((bits >> i) & 1) == 1
  379. if i < 8:
  380. self.modules[8][self.modules_count - i - 1] = mod
  381. elif i < 9:
  382. self.modules[8][15 - i - 1 + 1] = mod
  383. else:
  384. self.modules[8][15 - i - 1] = mod
  385. # fixed module
  386. self.modules[self.modules_count - 8][8] = not test
  387. def map_data(self, data, mask_pattern):
  388. inc = -1
  389. row = self.modules_count - 1
  390. bitIndex = 7
  391. byteIndex = 0
  392. mask_func = util.mask_func(mask_pattern)
  393. data_len = len(data)
  394. for col in range(self.modules_count - 1, 0, -2):
  395. if col <= 6:
  396. col -= 1
  397. col_range = (col, col - 1)
  398. while True:
  399. for c in col_range:
  400. if self.modules[row][c] is None:
  401. dark = False
  402. if byteIndex < data_len:
  403. dark = ((data[byteIndex] >> bitIndex) & 1) == 1
  404. if mask_func(row, c):
  405. dark = not dark
  406. self.modules[row][c] = dark
  407. bitIndex -= 1
  408. if bitIndex == -1:
  409. byteIndex += 1
  410. bitIndex = 7
  411. row += inc
  412. if row < 0 or self.modules_count <= row:
  413. row -= inc
  414. inc = -inc
  415. break
  416. def get_matrix(self):
  417. """
  418. Return the QR Code as a multidimensional array, including the border.
  419. To return the array without a border, set ``self.border`` to 0 first.
  420. """
  421. if self.data_cache is None:
  422. self.make()
  423. if not self.border:
  424. return self.modules
  425. width = len(self.modules) + self.border * 2
  426. code = [[False] * width] * self.border
  427. x_border = [False] * self.border
  428. for module in self.modules:
  429. code.append(x_border + cast(list[bool], module) + x_border)
  430. code += [[False] * width] * self.border
  431. return code
  432. def active_with_neighbors(self, row: int, col: int) -> ActiveWithNeighbors:
  433. context: list[bool] = []
  434. for r in range(row - 1, row + 2):
  435. for c in range(col - 1, col + 2):
  436. context.append(self.is_constrained(r, c) and bool(self.modules[r][c]))
  437. return ActiveWithNeighbors(*context)