kicad-qrgen/qrgen_action.py

from dataclasses import dataclass
from enum import Enum

import pcbnew
from pcbnew import VECTOR2I

import os
import qrcode
import wx


def relpath(p):
    return os.path.join(os.path.dirname(__file__), f"./{p}")

    
class QrSpec:
    @dataclass
    class Params:
        class Shape(Enum):
            ROUND = "round"
            SQUARE = "square"  # todo: squares / rect
                
        dot_size: float
        border_size: int
        invert: bool
        shape: Shape
        
        @classmethod
        def default(cls):
            return cls(dot_size=0.5, border_size=2, invert=False, shape=cls.Shape.ROUND)
        
        @classmethod
        def from_header(cls, header):
            _, h = header.split("qrgen!")
            opts = h.strip().split(" ")
            params = cls.default()
            for opt in opts:
                if "=" in opt:
                    inst, value = opt.split("=")
                else:
                    inst = opt
                    value = None
                
                match inst:
                    case "dot_size":
                        params.dot_size = float(value)
                    case "border_size":
                        params.border_size = int(value)
                    case "invert":
                        params.invert = True
                    case "shape":
                        params.shape = cls.Shape(value)
            
            return params
        
    def __init__(self, raw_text: str):
        self._raw = raw_text
        self.header = self.content = self.params = None
        self.init_success = False
        self.init()
        
    def init(self):
        self.header, self.content = self._raw.split("\n", maxsplit=1)
        self.params = self.Params.from_header(self.header)
        self.init_success = True
        
    def __repr__(self):
        return f"QrSpec(header={repr(self.header)}, params={repr(self.params)}, content={repr(self.content)})"
    
    def as_line_segments(self):
        q = qrcode.QRCode()
        q.add_data(self.content)
        q.border = self.params.border_size
        data = q.get_matrix()
        if self.params.invert:
            data = [[not p for p in row] for row in data]
        
        return array_to_lines(data)
        

def array_to_lines(array: list[list[bool]]) -> tuple[list[list[tuple[int, int]]], list[list[tuple[int, int]]]]:
    horiz = [line_to_segments(row) for row in array]
    vert = [line_to_segments(col[::-1]) for col in list(zip(*array[::-1]))]  # ?????
    return horiz, vert


def line_to_segments(line: list[bool]) -> list[tuple[int, int]]:
    in_seg = line[0]
    out = []
    start = 0

    for i, pix in enumerate(line):
        if pix and not in_seg:
            start = i
            in_seg = True
        elif not pix and in_seg:
            out.append((start, i-1))
            in_seg = False
    
    if in_seg:
        out.append((start, len(line)-1))

    return out
    

def handle_exception(e):
    dlg = wx.MessageDialog(None, f"error in processing: {e}")
    dlg.ShowModal()
    dlg.Destroy()
    return


def add_line(parent, start, end, width, layer):
    seg = pcbnew.PCB_SHAPE(parent)
    seg.SetShape(pcbnew.SHAPE_T_SEGMENT)
    seg.SetStart(start)
    seg.SetEnd(end)
    seg.SetWidth(int(width))
    seg.SetLayer(layer)
    parent.Add(seg)


def qrgen(source_elem: pcbnew.PCB_TEXT):
    raw = str(source_elem.GetText())
    layer = source_elem.GetLayer()
    x = source_elem.GetX()
    y = source_elem.GetY()
    qr = QrSpec(raw)
    board = pcbnew.GetBoard()
    
    h_lines, v_lines = qr.as_line_segments()
    params = qr.params
    dsize = int(params.dot_size * 1e6)
    
    all_segs = []

    for i, line in enumerate(h_lines):
        y_offset = i * dsize
        for start, end in line:
            all_segs.append((VECTOR2I(x + start * dsize, y + y_offset),
                             VECTOR2I(x + end * dsize, y + y_offset)))
            
    for i, col in enumerate(v_lines):
        x_offset = i * dsize
        for start, end in col:
            all_segs.append((VECTOR2I(x + x_offset, y + start * dsize),
                             VECTOR2I(x + x_offset, y + end * dsize)))
    
    for a, b in all_segs:
        add_line(board, a, b, dsize, layer)
    

class QrgenPluginAction(pcbnew.ActionPlugin):
    def defaults(self):
        self.name = "generate qr codes"
        self.category = "Modify PCB"
        self.description = "generate qr codes without leaving pcbnew"
        self.show_toolbar_button = True
        self.icon_file_name = relpath("qrgen.png")
    
    def Run(self):
        qr_text: list[pcbnew.PCB_TEXT] = [text for text in pcbnew.GetCurrentSelection()
                                          if isinstance(text, pcbnew.PCB_TEXT)
                                          and str(text.GetText()).startswith("qrgen!")]
        
        for elem in qr_text:
            try:
                qrgen(elem)
            except Exception as e:
                handle_exception(e)
        
        pcbnew.Refresh()