k2t/k2t.py

"""
Functions from converting between KLayout layer properties files and a custom TOML format
"""
from typing import Self, Type, Any
import tomllib
from dataclasses import dataclass, fields
from xml.etree import ElementTree


# Used when value is `None`, except on groups where they're left unset
DEFAULT_LINESTYLE = 'I0'
DEFAULT_DITHER = 'I1'

INDENT = '  '


def k2t(
        kfile: str,
        tfile: str,
        fill_name: bool = False,
        ) -> None:
    """
    Convert from KLayout layer properties xml (.lyp) to TOML

    Args:
        kfile: Path to klayout lyp file (input)
        tfile: Path for TOML file (output)
    """
    tree = ElementTree.parse(kfile)
    root = tree.getroot()

    assert root.tag == 'layer-properties'
    name = None
    entries = []
    patterns = []
    linestyles = []
    for child in root:
        if child.tag == 'custom-dither-pattern':
            patterns.append(Pattern.from_xml(child))
        elif child.tag == 'custom-line-style':
            linestyles.append(LineStyle.from_xml(child))
        elif child.tag == 'properties':
            entries.append(Entry.from_xml(child))
        elif child.tag == 'name':
            name = child.text
        else:
            raise Exception(f'Unexpected tag: {child.tag}')

    if name is not None:
        print(f'Discarding tab name: {name}')

    def add_names(entrs: list[Entry]) -> None:
        for entry in entrs:
            if entry.members:
                add_names(entry.members)

            if entry.name or not entry.source:
                continue
            entry.name = entry.source

    if fill_name:
        add_names(entries)

    with open(tfile, 'wt') as ff:
        for entry in entries:
            ff.write(entry.to_toml())
        for pattern in patterns:
            ff.write(pattern.to_toml())
        for linestyle in linestyles:
            ff.write(linestyle.to_toml())


def t2k(tfile: str, kfile: str) -> None:
    """
    Convert from TOML to KLayout layer properties xml (.lyp)

    Args:
        tfile: Path to TOML file (input)
        kfile: Path for klayout lyp file (output)
    """
    root = ElementTree.Element('layer-properties')

    with open(tfile, 'rb') as ff:
        ttree = tomllib.load(ff)

    ttree.setdefault('layers', [])
    ttree.setdefault('patterns', [])
    ttree.setdefault('linestyles', [])

    entries = [Entry.from_toml(entry)
               for entry in ttree['layers']]
    patterns = [Pattern.from_toml(pat)
                for pat in ttree['patterns']]
    linestyles = [LineStyle.from_toml(ls)
                  for ls in ttree['linestyles']]

    for entry in entries:
        root.append(entry.to_xml())
    for pattern in patterns:
        root.append(pattern.to_xml())
    for linestyle in linestyles:
        root.append(linestyle.to_xml())

    ktree = ElementTree.ElementTree(root)
    ElementTree.indent(ktree)
    ktree.write(kfile)


@dataclass(slots=True)
class Entry:
    name: str | None = None
    source: str | None = None
    frame_color: str | None = None
    fill_color: str | None = None

    members: list | None = None
    expanded: bool = False

    dither_pattern: str | None = None
    line_style: str | None = None
    frame_brightness: int = 0
    fill_brightness: int = 0
    width: int = 1
    animation: int = 0

    valid: bool = True
    visible: bool = True
    transparent: bool = False
    marked: bool = False
    xfill: bool = False

    @classmethod
    def from_xml(cls: Type[Self], el: ElementTree.Element) -> Self:
        assert el.tag in ('properties', 'group-members')

        members = []
        args: dict[str, Any] = {}
        for child in el:
            if not child.text:
                continue

            tag = child.tag
            dtag = tag.replace('-', '_')
            if tag in ('dither-pattern', 'line-style', 'source', 'name', 'frame-color', 'fill-color'):
                args[dtag] = child.text
            elif tag in ('valid', 'visible', 'transparent', 'marked', 'xfill', 'expanded'):
                val = {
                    'false': False,
                    'true': True,
                    }[child.text.lower()]
                args[tag] = val
            elif tag in ('frame-brightness', 'fill-brightness', 'width', 'animation'):
                args[dtag] = int(child.text)
            elif tag == 'group-members':
                members.append(Entry.from_xml(child))
            else:
                raise Exception(f'Unexpected tag: {child.tag}')

        if not members:
            members = None
            args.setdefault('line_style', DEFAULT_LINESTYLE)
            args.setdefault('dither_pattern', DEFAULT_DITHER)

        return cls(members=members, **args)

    def to_xml(self, is_member: bool = False) -> ElementTree.Element:
        main_tag = 'group-members' if is_member else 'properties'
        el = ElementTree.Element(main_tag)
        for field in fields(self):
            val = getattr(self, field.name)

            if field.name == 'members':
                continue

            tag = field.name.replace('_', '-')
            el2 = ElementTree.Element(tag)
            if val is not None:
                if not isinstance(val, str):
                    val = str(val).lower()
                el2.text = val
            elif tag == 'source':
                el2.text = '*/*@*'
            el.append(el2)

        if self.members:
            for member in self.members:
                el.append(member.to_xml(is_member=True))

        return el

    def to_toml(self, depth: int = 0) -> str:
        indent = INDENT * depth
        parts = ('layers',) + ('members',) * depth

        s = indent + '[[' + '.'.join(parts)  + ']]\n'
        for field in fields(self):
            val = getattr(self, field.name)
            if val == field.default:
                continue
            if field.name == 'members':
                continue

            if isinstance(val, str):
                vstr = '"' + toml_escape(val) + '"'
            elif isinstance(val, bool):
                vstr = str(val).lower()
            elif isinstance(val, int):
                vstr = str(val)

            key = field.name.replace('_', '-')
            s += f'{indent}{key} = {vstr}\n'

        s += '\n'

        if self.members:
            for member in self.members:
                s += member.to_toml(depth + 1)

        return s

    @classmethod
    def from_toml(cls: Type[Self], tree: dict) -> Self:
        members = []
        args: dict[str, Any] = {}
        for key, val in tree.items():
            if key != 'members':
                args[key.replace('-', '_')] = val

        if 'members' in tree:
            for child in tree['members']:
                members.append(Entry.from_toml(child))

        if not members:
            members = None
            args.setdefault('line_style', DEFAULT_LINESTYLE)
            args.setdefault('dither_pattern', DEFAULT_DITHER)

        return cls(members=members, **args)


@dataclass(slots=True)
class Pattern:
    pattern: str              # 2D,  Nx8 or Nx16 or Nx32
    name: str | None = None
    order: int | None = None

    @classmethod
    def from_xml(cls: Type[Self], el: ElementTree.Element) -> Self:
        assert el.tag == 'custom-dither-pattern'

        lines = []
        args: dict[str, Any] = {}
        for child in el:
            if not child.text:
                continue

            tag = child.tag
            if tag == 'name':
                args[tag] = child.text
            elif tag == 'order':
                args[tag] = int(child.text)
            elif tag == 'pattern':
                for line in child:
                    assert line.tag == 'line'
                    assert line.text is not None
                    lines.append(line.text)
            else:
                raise Exception(f'Unexpected tag: {tag}')

        if not lines:
            raise Exception('No pattern found in custom-dither-pattern!')
        pattern = '\n'.join(lines)

        return cls(pattern=pattern, **args)

    def to_xml(self) -> ElementTree.Element:
        el = ElementTree.Element('custom-dither-pattern')
        if self.name is not None:
            el2 = ElementTree.Element('name')
            el2.text = self.name
            el.append(el2)
        if self.order is not None:
            el2 = ElementTree.Element('order')
            el2.text = str(self.order)
            el.append(el2)

        pat = ElementTree.Element('pattern')
        for line in self.pattern.strip().split('\n'):
            el2 = ElementTree.Element('line')
            el2.text = line
            pat.append(el2)
        el.append(pat)

        return el

    def to_toml(self) -> str:
        s = '[[patterns]]\n'
        if self.name is not None:
            s += f'name = "{toml_escape(self.name)}"\n'
        if self.order is not None:
            s += f'order = {self.order}\n'

        s += 'pattern = """\n'
        s += self.pattern + '"""\n\n'

        return s

    @classmethod
    def from_toml(cls: Type[Self], tree: dict) -> Self:
        lines = tree['pattern'].split('\n')
        n_px = len(lines[0])
        assert all(len(line) == n_px for line in lines[1:])
        return cls(**tree)


@dataclass(slots=True)
class LineStyle:
    pattern: str               # 1D, up to 32 values
    name: str | None = None
    order: int | None = None

    @classmethod
    def from_xml(cls: Type[Self], el: ElementTree.Element) -> Self:
        assert el.tag == 'custom-line-style'

        args: dict[str, Any] = {}
        for child in el:
            if not child.text:
                continue

            tag = child.tag
            if tag in ('name', 'pattern'):
                args[tag] = child.text
            elif tag == 'order':
                args[tag] = int(child.text)
            else:
                raise Exception(f'Unexpected tag: {tag}')
        return cls(**args)

    def to_xml(self) -> ElementTree.Element:
        el = ElementTree.Element('custom-line-style')
        if self.name is not None:
            el2 = ElementTree.Element('name')
            el2.text = self.name
            el.append(el2)
        if self.order is not None:
            el2 = ElementTree.Element('order')
            el2.text = str(self.order)
            el.append(el2)
        el2 = ElementTree.Element('pattern')
        el2.text = self.pattern
        el.append(el2)
        return el

    def to_toml(self) -> str:
        s = '[[linestyles]]\n'
        if self.name is not None:
            s += f'name = "{toml_escape(self.name)}"\n'
        if self.order is not None:
            s += f'order = {self.order}\n'
        s += f'pattern = "{self.pattern}"\n\n'
        return s

    @classmethod
    def from_toml(cls: Type[Self], tree: dict) -> Self:
        assert len(tree['pattern']) <= 32
        return cls(**tree)


def toml_escape(string: str) -> str:
    return string.replace('\\', '\\\\').replace('"', '\\"')