krita_grid_generator/pykrita/grid_generator/grid_docker.py

import math

from krita import DockWidget, DockWidgetFactory, DockWidgetFactoryBase, Krita
from PyQt5.QtCore import QPointF
from PyQt5.QtWidgets import (
    QCheckBox,
    QComboBox,
    QDoubleSpinBox,
    QHBoxLayout,
    QLabel,
    QPushButton,
    QVBoxLayout,
    QWidget,
)

DOCKER_ID = "grid_generator_docker"


class GridDocker(DockWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Grid Generator")
        self.setup_ui()

    def setup_ui(self):
        main_widget = QWidget()
        layout = QVBoxLayout(main_widget)

        # Thickness
        thickness_layout = QHBoxLayout()
        thickness_layout.addWidget(QLabel("Thickness (px):"))
        self.thickness_spin = QDoubleSpinBox()
        self.thickness_spin.setRange(0.1, 1000.0)
        self.thickness_spin.setValue(2.0)
        thickness_layout.addWidget(self.thickness_spin)
        layout.addLayout(thickness_layout)

        # Width
        self.width_container = QWidget()
        width_layout = QHBoxLayout(self.width_container)
        width_layout.setContentsMargins(0, 0, 0, 0)
        width_layout.addWidget(QLabel("Width (px):"))
        self.width_spin = QDoubleSpinBox()
        self.width_spin.setRange(1.0, 10000.0)
        self.width_spin.setValue(100.0)
        self.width_spin.valueChanged.connect(self.on_width_changed)
        width_layout.addWidget(self.width_spin)
        layout.addWidget(self.width_container)

        # Height
        self.height_container = QWidget()
        height_layout = QHBoxLayout(self.height_container)
        height_layout.setContentsMargins(0, 0, 0, 0)
        height_layout.addWidget(QLabel("Height (px):"))
        self.height_spin = QDoubleSpinBox()
        self.height_spin.setRange(1.0, 10000.0)
        self.height_spin.setValue(100.0)
        self.height_spin.valueChanged.connect(self.on_height_changed)
        height_layout.addWidget(self.height_spin)
        layout.addWidget(self.height_container)

        # Side Length (Hexagon Only)
        self.side_container = QWidget()
        side_layout = QHBoxLayout(self.side_container)
        side_layout.setContentsMargins(0, 0, 0, 0)
        side_layout.addWidget(QLabel("Side Length (px):"))
        self.side_spin = QDoubleSpinBox()
        self.side_spin.setRange(1.0, 10000.0)
        self.side_spin.setValue(50.0)
        side_layout.addWidget(self.side_spin)
        layout.addWidget(self.side_container)
        self.side_container.hide()

        # Lock Width/Height
        self.lock_check = QCheckBox("Lock Aspect Ratio")
        self.lock_check.setChecked(True)
        layout.addWidget(self.lock_check)

        # Grid Type
        type_layout = QHBoxLayout()
        type_layout.addWidget(QLabel("Type:"))
        self.type_combo = QComboBox()
        self.type_combo.addItems(["Square", "Hexagon"])
        self.type_combo.currentTextChanged.connect(self.on_type_changed)
        type_layout.addWidget(self.type_combo)
        layout.addLayout(type_layout)

        # Color Option
        color_layout = QHBoxLayout()
        color_layout.addWidget(QLabel("Color:"))
        self.color_combo = QComboBox()
        self.color_combo.addItems(["Foreground", "Background"])
        color_layout.addWidget(self.color_combo)
        layout.addLayout(color_layout)

        # Buttons
        self.ok_button = QPushButton("Generate")
        self.ok_button.clicked.connect(self.generate)
        layout.addWidget(self.ok_button)

        layout.addStretch()
        self.setWidget(main_widget)

    def on_type_changed(self, grid_type):
        is_square = grid_type == "Square"
        self.width_container.setVisible(is_square)
        self.height_container.setVisible(is_square)
        self.lock_check.setVisible(is_square)
        self.side_container.setVisible(not is_square)

    def on_width_changed(self, value):
        if self.lock_check.isChecked():
            self.height_spin.blockSignals(True)
            self.height_spin.setValue(value)
            self.height_spin.blockSignals(False)

    def on_height_changed(self, value):
        if self.lock_check.isChecked():
            self.width_spin.blockSignals(True)
            self.width_spin.setValue(value)
            self.width_spin.blockSignals(False)

    def canvasChanged(self, canvas):
        pass

    def get_color_hex(self):
        app = Krita.instance()
        window = app.activeWindow()
        if not window:
            return "#000000"

        view = window.activeView()
        if not view:
            return "#000000"

        color = None
        if self.color_combo.currentText() == "Foreground":
            color = view.foregroundColor()
        else:
            color = view.backgroundColor()

        color.setColorSpace("RGBA", "U8", "sRGB built-in")
        comps = color.componentsOrdered()

        r = int(max(0, min(1, comps[0])) * 255)
        g = int(max(0, min(1, comps[1])) * 255)
        b = int(max(0, min(1, comps[2])) * 255)

        return f"#{r:02x}{g:02x}{b:02x}"

    def generate(self):
        app = Krita.instance()
        doc = app.activeDocument()
        if not doc:
            return

        thickness = self.thickness_spin.value()
        grid_type = self.type_combo.currentText()
        color_hex = self.get_color_hex()

        # Always create a new vector layer for the grid
        target_node = doc.createVectorLayer("grid")
        doc.rootNode().addChildNode(target_node, None)

        if hasattr(target_node, "addShapesFromSvg"):
            if grid_type == "Square":
                width = self.width_spin.value()
                height = self.height_spin.value()
                self.draw_square_grid(
                    target_node, doc, thickness, width, height, color_hex
                )
            else:
                side_length = self.side_spin.value()
                self.draw_hexagon_grid(
                    target_node, doc, thickness, side_length, color_hex
                )
        else:
            print(
                "Error: Target node does not support SVG shapes. Please ensure you are using Krita 5.0+."
            )

        doc.refreshProjection()

    def draw_square_grid(self, node, doc, thickness, width, height, color_hex):
        svg = f'<svg><g fill="none" stroke="{color_hex}" stroke-width="{thickness}">'
        x = 0
        while x <= doc.width():
            svg += f'<line x1="{x}" y1="0" x2="{x}" y2="{doc.height()}" />'
            x += width
        y = 0
        while y <= doc.height():
            svg += f'<line x1="0" y1="{y}" x2="{doc.width()}" y2="{y}" />'
            y += height
        svg += "</g></svg>"
        node.addShapesFromSvg(svg)

    def draw_hexagon_grid(self, node, doc, thickness, side_length, color_hex):
        h_width = math.sqrt(3) * side_length
        v_spacing = 1.5 * side_length

        svg = f'<svg><g fill="none" stroke="{color_hex}" stroke-width="{thickness}">'

        rows = int(doc.height() / v_spacing) + 2
        cols = int(doc.width() / h_width) + 2

        for r in range(rows):
            offset_x = (h_width / 2) if (r % 2 == 1) else 0
            for c in range(cols):
                cx = c * h_width + offset_x
                cy = r * v_spacing

                v = []
                for i in range(6):
                    angle_deg = 60 * i - 90
                    angle_rad = math.pi / 180 * angle_deg
                    px = cx + side_length * math.cos(angle_rad)
                    py = cy + side_length * math.sin(angle_rad)
                    v.append((px, py))

                points = " ".join([f"{px},{py}" for px, py in v])
                svg += f'<polygon points="{points}" />'

        svg += "</g></svg>"
        node.addShapesFromSvg(svg)


Krita.instance().addDockWidgetFactory(
    DockWidgetFactory(DOCKER_ID, DockWidgetFactoryBase.DockRight, GridDocker)
)