base.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. import abc
  2. from typing import TYPE_CHECKING, Any, Optional, Union
  3. from qrcode.image.styles.moduledrawers.base import QRModuleDrawer
  4. if TYPE_CHECKING:
  5. from qrcode.main import ActiveWithNeighbors, QRCode
  6. DrawerAliases = dict[str, tuple[type[QRModuleDrawer], dict[str, Any]]]
  7. class BaseImage:
  8. """
  9. Base QRCode image output class.
  10. """
  11. kind: Optional[str] = None
  12. allowed_kinds: Optional[tuple[str]] = None
  13. needs_context = False
  14. needs_processing = False
  15. needs_drawrect = True
  16. def __init__(self, border, width, box_size, *args, **kwargs):
  17. self.border = border
  18. self.width = width
  19. self.box_size = box_size
  20. self.pixel_size = (self.width + self.border * 2) * self.box_size
  21. self.modules = kwargs.pop("qrcode_modules")
  22. self._img = self.new_image(**kwargs)
  23. self.init_new_image()
  24. @abc.abstractmethod
  25. def drawrect(self, row, col):
  26. """
  27. Draw a single rectangle of the QR code.
  28. """
  29. def drawrect_context(self, row: int, col: int, qr: "QRCode"):
  30. """
  31. Draw a single rectangle of the QR code given the surrounding context
  32. """
  33. raise NotImplementedError("BaseImage.drawrect_context") # pragma: no cover
  34. def process(self):
  35. """
  36. Processes QR code after completion
  37. """
  38. raise NotImplementedError("BaseImage.drawimage") # pragma: no cover
  39. @abc.abstractmethod
  40. def save(self, stream, kind=None):
  41. """
  42. Save the image file.
  43. """
  44. def pixel_box(self, row, col):
  45. """
  46. A helper method for pixel-based image generators that specifies the
  47. four pixel coordinates for a single rect.
  48. """
  49. x = (col + self.border) * self.box_size
  50. y = (row + self.border) * self.box_size
  51. return (
  52. (x, y),
  53. (x + self.box_size - 1, y + self.box_size - 1),
  54. )
  55. @abc.abstractmethod
  56. def new_image(self, **kwargs) -> Any:
  57. """
  58. Build the image class. Subclasses should return the class created.
  59. """
  60. def init_new_image(self):
  61. pass
  62. def get_image(self, **kwargs):
  63. """
  64. Return the image class for further processing.
  65. """
  66. return self._img
  67. def check_kind(self, kind, transform=None):
  68. """
  69. Get the image type.
  70. """
  71. if kind is None:
  72. kind = self.kind
  73. allowed = not self.allowed_kinds or kind in self.allowed_kinds
  74. if transform:
  75. kind = transform(kind)
  76. if not allowed:
  77. allowed = kind in self.allowed_kinds
  78. if not allowed:
  79. raise ValueError(f"Cannot set {type(self).__name__} type to {kind}")
  80. return kind
  81. def is_eye(self, row: int, col: int):
  82. """
  83. Find whether the referenced module is in an eye.
  84. """
  85. return (
  86. (row < 7 and col < 7)
  87. or (row < 7 and self.width - col < 8)
  88. or (self.width - row < 8 and col < 7)
  89. )
  90. class BaseImageWithDrawer(BaseImage):
  91. default_drawer_class: type[QRModuleDrawer]
  92. drawer_aliases: DrawerAliases = {}
  93. def get_default_module_drawer(self) -> QRModuleDrawer:
  94. return self.default_drawer_class()
  95. def get_default_eye_drawer(self) -> QRModuleDrawer:
  96. return self.default_drawer_class()
  97. needs_context = True
  98. module_drawer: "QRModuleDrawer"
  99. eye_drawer: "QRModuleDrawer"
  100. def __init__(
  101. self,
  102. *args,
  103. module_drawer: Union[QRModuleDrawer, str, None] = None,
  104. eye_drawer: Union[QRModuleDrawer, str, None] = None,
  105. **kwargs,
  106. ):
  107. self.module_drawer = (
  108. self.get_drawer(module_drawer) or self.get_default_module_drawer()
  109. )
  110. # The eye drawer can be overridden by another module drawer as well,
  111. # but you have to be more careful with these in order to make the QR
  112. # code still parseable
  113. self.eye_drawer = self.get_drawer(eye_drawer) or self.get_default_eye_drawer()
  114. super().__init__(*args, **kwargs)
  115. def get_drawer(
  116. self, drawer: Union[QRModuleDrawer, str, None]
  117. ) -> Optional[QRModuleDrawer]:
  118. if not isinstance(drawer, str):
  119. return drawer
  120. drawer_cls, kwargs = self.drawer_aliases[drawer]
  121. return drawer_cls(**kwargs)
  122. def init_new_image(self):
  123. self.module_drawer.initialize(img=self)
  124. self.eye_drawer.initialize(img=self)
  125. return super().init_new_image()
  126. def drawrect_context(self, row: int, col: int, qr: "QRCode"):
  127. box = self.pixel_box(row, col)
  128. drawer = self.eye_drawer if self.is_eye(row, col) else self.module_drawer
  129. is_active: Union[bool, ActiveWithNeighbors] = (
  130. qr.active_with_neighbors(row, col)
  131. if drawer.needs_neighbors
  132. else bool(qr.modules[row][col])
  133. )
  134. drawer.drawrect(box, is_active)