pcbdraw.py

#!/usr/bin/env python

import argparse
import json
import math
import os
import re
import shutil
import sys
import tempfile

import pcbnew
from lxml import etree

default_style = {
    "copper": "#417e5a",
    "board": "#4ca06c",
    "silk": "#f0f0f0",
    "pads": "#b5ae30",
    "outline": "#000000"
}

class SvgPathItem:
    def __init__(self, path):
        path = re.sub(r"([MLA])(\d+)", r"\1 \2", path)
        path = re.split("[, ]", path)
        path = filter( lambda x: x, path)
        if path[0] != "M":
            raise SyntaxError("Only paths with absolute position are supported")
        self.start = tuple(map(float, path[1:3]))
        path = path[3:]
        if path[0] == "L":
            x = float(path[1])
            y = float(path[2])
            self.end = (x, y)
            self.type = path[0]
            self.args = None
        elif path[0] == "A":
            args = map(float, path[1:8])
            self.end = (args[5], args[6])
            self.args = args[0:5]
            self.type = path[0]
        else:
            raise SyntaxError("Unsupported path element " + path[0])

    @staticmethod
    def is_same(p1, p2):
        dx = p1[0] - p2[0]
        dy = p1[1] - p2[1]
        return math.sqrt(dx*dx+dy*dy) < 5

    def format(self, first):
        ret = ""
        if first:
            ret += " M {} {} ".format(*self.start)
        ret += self.type
        if self.args:
            ret += " " + " ".join(map(lambda x: str(x).rstrip('0').rstrip('.'), self.args))
        ret += " {} {} ".format(*self.end)
        return ret

    def flip(self):
        self.start, self.end = self.end, self.start
        if self.type == "A":
            self.args[4] = 1 if self.args[4] < 0.5 else 0

def unique_prefix():
    unique_prefix.counter += 1
    return "pref_" + str(unique_prefix.counter)
unique_prefix.counter = 0

def ki2dmil(val):
    return val / 2540

def extract_svg_content(filename):
    prefix = unique_prefix() + "_"
    root = etree.parse(filename).getroot()
    # We have to ensure all Ids in SVG are unique. Let's make it nasty by
    # collecting all ids and doing search & replace
    # Potentially dangerous (can break user text)
    ids = []
    for el in root.getiterator():
        if "id" in el.attrib and el.attrib["id"] != "origin":
            ids.append(el.attrib["id"])
    with open(filename) as f:
        content = f.read()
    for i in ids:
        content = content.replace("#"+i, "#" + prefix + i)
    root = etree.fromstring(content)
    # Remove SVG namespace to ease our lifes and change ids
    for el in root.getiterator():
        if "id" in el.attrib and el.attrib["id"] != "origin":
            el.attrib["id"] = prefix + el.attrib["id"]
        if '}' in str(el.tag):
            el.tag = el.tag.split('}', 1)[1]
    return [ x for x in root if x.tag and x.tag not in ["title", "desc"]]

def strip_fill_svg(root):
    keys = ["fill", "stroke"]
    for el in root.getiterator():
        if "style" in el.attrib:
            s = el.attrib["style"].split(";")
            s = filter(lambda x: x.strip().split(":")[0] not in keys, s)
            el.attrib["style"] = ";".join(s).replace("  ", " ").strip()

def empty_svg(**attrs):
    document = etree.ElementTree(etree.fromstring(
        """<?xml version="1.0" standalone="no"?>
        <!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
            "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
        <svg xmlns="http://www.w3.org/2000/svg" version="1.1"
            width="29.7002cm" height="21.0007cm" viewBox="0 0 116930 82680 ">
            <title>Picture generated by pcb2svg</title>
            <desc>Picture generated by pcb2svg</desc>
        </svg>"""))
    root = document.getroot()
    for key, value in attrs.items():
        root.attrib[key] = value
    return document

def get_board_polygon(svg_elements):
    """
    Try to connect independents segments on Edge.Cuts and form a polygon
    return SVG path element with the polygon
    """
    elements = []
    path = ""
    for group in svg_elements:
        for svg_element in group:
            if svg_element.tag == "path":
                elements.append(SvgPathItem(svg_element.attrib["d"]))
            elif svg_element.tag == "circle":
                # Convert circle to path
                att = svg_element.attrib
                s = " M {0} {1} m-{2} 0 a {2} {2} 0 1 0 {3} 0 a {2} {2} 0 1 0 -{3} 0 ".format(
                    att["cx"], att["cy"], att["r"], 2 * float(att["r"]))
                path += s
    outline = [elements[0]]
    elements = elements[1:]
    while True:
        size = len(outline)
        for i, e in enumerate(elements):
            if SvgPathItem.is_same(outline[0].start, e.end):
                outline.insert(0, e)
            elif SvgPathItem.is_same(outline[0].start, e.start):
                e.flip()
                outline.insert(0, e)
            elif SvgPathItem.is_same(outline[-1].end, e.start):
                outline.append(e)
            elif SvgPathItem.is_same(outline[-1].end, e.end):
                e.flip()
                outline.append(e)
            else:
                continue
            del elements[i]
            break
        if size == len(outline):
            first = True
            for x in outline:
                path += x.format(first)
                first = False
            if elements:
                outline = [elements[0]]
                elements = elements[1:]
            else:
                e = etree.Element("path", d=path, style="fill-rule=evenodd;")
                return e

def process_board_substrate_layer(container, name, source, colors):
    layer = etree.SubElement(container, "g", id="substrate-"+name,
        style="fill:{0}; stroke:{0};".format(colors[name]))
    if name == "pads":
        layer.attrib["mask"] = "url(#pads-mask)";
    for element in extract_svg_content(source):
        strip_fill_svg(element)
        layer.append(element)

def process_board_substrate_base(container, name, source, colors):
    clipPath = etree.SubElement(etree.SubElement(container, "defs"), "clipPath")
    clipPath.attrib["id"] = "cut-off"
    clipPath.append(get_board_polygon(extract_svg_content(source)))

    layer = etree.SubElement(container, "g", id="substrate-"+name,
        style="fill:{0}; stroke:{0};".format(colors[name]))
    layer.append(get_board_polygon(extract_svg_content(source)))
    outline = etree.SubElement(layer, "g",
        style="fill:{0}; stroke: {0};".format(colors["outline"]))
    for element in extract_svg_content(source):
        strip_fill_svg(element)
        outline.append(element)

def process_board_substrate_mask(container, name, source, colors):
    mask = etree.SubElement(etree.SubElement(container, "defs"), "mask")
    mask.attrib["id"] = name
    for element in extract_svg_content(source):
        for item in element.getiterator():
            if "style" in item.attrib:
                # KiCAD plots in black, for mask we need white
                item.attrib["style"] = item.attrib["style"].replace("#000000", "#ffffff");
            mask.append(element)

def get_board_substrate(board, colors, holes):
    """
    Plots all front layers from the board and arranges them in a visually appealing style.
    return SVG g element with the board substrate
    """
    toPlot = [
        ("board", [pcbnew.Edge_Cuts], process_board_substrate_base),
        ("copper", [pcbnew.F_Cu], process_board_substrate_layer),
        ("pads", [pcbnew.F_Cu], process_board_substrate_layer),
        ("pads-mask", [pcbnew.F_Mask], process_board_substrate_mask),
        ("silk", [pcbnew.F_SilkS], process_board_substrate_layer),
        ("outline", [pcbnew.Edge_Cuts], process_board_substrate_layer)]
    container = etree.Element('g')
    container.attrib["clip-path"] = "url(#cut-off)";
    tmp = tempfile.mkdtemp()
    pctl = pcbnew.PLOT_CONTROLLER(board)
    popt = pctl.GetPlotOptions()
    popt.SetOutputDirectory(tmp)
    popt.SetScale(1)
    popt.SetMirror(False)
    try:
        popt.SetPlotOutlineMode(False)
    except:
        # Method does not exist in older versions of KiCad
        pass
    popt.SetTextMode(pcbnew.PLOTTEXTMODE_STROKE)
    for f, layers, _ in toPlot:
        pctl.OpenPlotfile(f, pcbnew.PLOT_FORMAT_SVG, f)
        for l in layers:
            pctl.SetColorMode(False)
            pctl.SetLayer(l)
            pctl.PlotLayer()
    pctl.ClosePlot()
    for f, _, process in toPlot:
        for svg_file in os.listdir(tmp):
            if svg_file.endswith("-" + f + ".svg"):
                process(container, f, os.path.join(tmp, svg_file), colors)
    shutil.rmtree(tmp)

    if holes:
        container.append(get_hole_mask(board))
        container.attrib["mask"] = "url(#hole-mask)";
    return container

def walk_components(board, export):
    module = board.GetModules()
    while True:
        if not module:
            return
        # Top is for Eagle boards imported to KiCAD
        if str(module.GetLayerName()) not in ["Top", "F.Cu"]:
            module = module.Next()
            continue
        lib = str(module.GetFPID().GetLibNickname()).strip()
        try:
            name = str(module.GetFPID().GetFootprintName()).strip()
        except AttributeError:
            # it seems we are working on Kicad >4.0.6, which has a changed method name
            name = str(module.GetFPID().GetLibItemName()).strip()
        value = unicode(module.GetValue()).strip()
        ref = unicode(module.GetReference()).strip()
        center = module.GetCenter()
        orient = math.radians(module.GetOrientation() / 10)
        pos = (center.x, center.y, orient)
        export(lib, name, value, ref, pos)
        module = module.Next()

def get_hole_mask(board):
    defs = etree.Element("defs")
    mask = etree.SubElement(defs, "mask", id="hole-mask")
    container = etree.SubElement(mask, "g")

    bb = board.ComputeBoundingBox();
    bg = etree.SubElement(container, "rect", x="0", y="0", fill="white")
    bg.attrib["x"] = str(ki2dmil(bb.GetX()))
    bg.attrib["y"] = str(ki2dmil(bb.GetY()))
    bg.attrib["width"] = str(ki2dmil(bb.GetWidth()))
    bg.attrib["height"] = str(ki2dmil(bb.GetHeight()))

    module = board.GetModules()
    while module:
        if module.GetPadCount() == 0:
            module = module.Next()
            continue
        pad = module.Pads()
        try:
            pad.GetPosition()
        except:
            # Newest nightly renames Pads to PadList
            pad = module.PadsList()
        orient = module.GetOrientation()
        while pad:
            pos = pad.GetPosition()
            pos.x = ki2dmil(pos.x)
            pos.y = ki2dmil(pos.y)
            size = map(ki2dmil, pad.GetDrillSize())
            if size[0] > 0 and size[1] > 0:
                if size[0] < size[1]:
                    stroke = size[0]
                    length = size[1] - size[0]
                    points = "{} {} {} {}".format(0, -length / 2, 0, length / 2)
                else:
                    stroke = size[1]
                    length = size[0] - size[1]
                    points = "{} {} {} {}".format(-length / 2, 0, length / 2, 0)
                el = etree.SubElement(container, "polyline")
                el.attrib["stroke-linecap"] = "round"
                el.attrib["stroke"] = "black"
                el.attrib["stroke-width"] = str(stroke)
                el.attrib["points"] = points
                el.attrib["transform"] = "translate({} {}) rotate({})".format(
                    pos.x, pos.y, -orient)
            pad = pad.Next()
        module = module.Next()
    return defs

def get_model_file(paths, lib, name, ref, remapping):
    """ Find model file in library considering component remapping """
    for path in paths:
        if ref in remapping:
            lib, name = tuple(remapping[ref].split(":"))
        f = os.path.join(path, lib, name + ".svg")
        if os.path.isfile(f):
            return f
    return None

def print_component(paths, lib, name, value, ref, pos, remapping={}):
    f = get_model_file(paths, lib, name, ref, remapping)
    msg = "{} with package {}:{} at [{},{},{}] -> {}".format(
        ref, lib, name, pos[0], pos[1], math.degrees(pos[2]), f if f else "Not found")
    print(msg)

def component_from_library(parent, paths, lib, name, value, ref, pos, placeholder=True, remapping={}):
    if not name:
        return
    f = get_model_file(paths, lib, name, ref, remapping)
    if not f:
        print("Warning: component '{}' from library '{}' was not found".format(name, lib))
        if placeholder:
            etree.SubElement(parent, "rect", x=str(ki2dmil(pos[0]) - 150), y=str(ki2dmil(pos[1]) - 150),
                             width="300", height="300", style="fill:red;")
        return
    parent.append(etree.Comment("{}:{}".format(lib, name)))
    r = etree.SubElement(parent, "g")
    for x in extract_svg_content(f):
        r.append(x)
    origin_x = 0
    origin_y = 0
    origin = r.find(".//*[@id='origin']")
    if origin is not None:
        origin_x = float(origin.attrib["x"])
        origin_y = float(origin.attrib["y"])
        origin.getparent().remove(origin)
    else:
        print("Warning: component '{}' from library '{}' has no ORIGIN".format(name, lib))
    r.attrib["transform"] = "translate({} {}) scale(393.700787402) rotate({}) translate({}, {})".format(
            ki2dmil(pos[0]), ki2dmil(pos[1]),
            -math.degrees(pos[2]), -origin_x, -origin_y)

def load_style(style_file):
    try:
        with open(style_file, "r") as f:
            style = json.load(f)
    except IOError:
        raise RuntimeError("Cannot open style " + style_file)
    required = set(["copper", "board", "silk", "pads", "outline"])
    missing = required - set(style.keys())
    if missing:
        raise RuntimeError("Missing following keys in style {}: {}"
                           .format(style_file, ", ".join(missing)))
    extra = set(style.keys()) - required
    for x in extra:
        print("Warning: extra key '" + x + "' in style")
    # ToDo: Check validity of colors (SVG compatible format)
    return style

def load_remapping(remap_file):
    if not remap_file:
        return {}
    try:
        with open(remap_file, "r") as f:
            return json.load(f)
    except IOError:
        raise RuntimeError("Cannot open remapping file " + remap_file)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--style", help="JSON file with board style")
    parser.add_argument("libraries", help="directories containing SVG footprints")
    parser.add_argument("output", help="destination for final SVG")
    parser.add_argument("board", help=".kicad_pcb file to draw")
    parser.add_argument("-p", "--placeholder", action="store_true",
                        help="show placeholder for missing components")
    parser.add_argument("-m", "--remap",
                        help="JSON file with map part reference to <lib>:<model> to remap packages")
    parser.add_argument("-l", "--list-components", action="store_true",
                        help="Dry run, just list the components")
    parser.add_argument("--no-drillholes", action="store_true", help="Do not make holes transparent")

    args = parser.parse_args()
    args.libraries = args.libraries.split(',')
    print(args)

    try:
        if args.style:
            style = load_style(args.style)
        else:
            style = default_style
        remapping = load_remapping(args.remap)
    except RuntimeError as e:
        print(e.message)
        sys.exit(1)
    try:
        print("Please ignore following debug output of KiCAD Python API")
        board = pcbnew.LoadBoard(args.board)
        print("End of KiCAD debug output")
    except IOError:
        print("Cannot open board " + args.board)
        sys.exit(1)

    if args.list_components:
        walk_components(board, lambda lib, name, val, ref, pos:
                        print_component(args.libraries, lib, name, val, ref, pos,
                                        remapping=remapping))
        sys.exit(0)

    bb = board.ComputeBoundingBox()
    document = empty_svg(
        width="{}cm".format(bb.GetWidth()/10000000.0),
        height="{}cm".format(bb.GetHeight()/10000000.0),
        viewBox="0 0 {} {}".format(ki2dmil(bb.GetWidth()), ki2dmil(bb.GetHeight())))
    wrapper = etree.SubElement(document.getroot(), "g",
        transform="translate({}, {})".format(ki2dmil(-bb.GetX()), ki2dmil(-bb.GetY())))

    wrapper.append(get_board_substrate(board, style, not args.no_drillholes))
    walk_components(board, lambda lib, name, val, ref, pos:
        component_from_library(wrapper, args.libraries, lib, name, val, ref, pos,
                               placeholder=args.placeholder, remapping=remapping))
    document.write(args.output)