masque/masque/file/gdsii/writer.py

"""
GDSII writer for eager and source-backed libraries.

The generic mutable overlay and ports-importing view live in `masque.library`.
This module preserves source-backed GDS copy-through behavior where possible,
falling back to normal pattern serialization when a cell has been materialized
or remapped.
"""
from __future__ import annotations

from typing import IO, TYPE_CHECKING, Any, Protocol, cast, runtime_checkable
import gzip
import logging
import pathlib

from . import klamath
from ..utils import tmpfile
from ...error import LibraryError
from ...library import IBorrowing, ILibraryView, IMaterializable

if TYPE_CHECKING:
    from collections.abc import Mapping

    from ...pattern import Pattern


logger = logging.getLogger(__name__)


@runtime_checkable
class _GdsInfoSource(Protocol):
    """Structural capability for propagating GDS header metadata."""
    library_info: dict[str, Any]


@runtime_checkable
class _GdsRawCellSource(Protocol):
    """GDS-specific raw-structure copy-through capability."""

    def can_copy_raw_struct(self, name: str) -> bool: ...

    def raw_struct_bytes(self, name: str) -> bytes: ...


def _resolve_raw_struct(
        library: ILibraryView,
        name: str,
        ) -> tuple[_GdsRawCellSource, str] | None:
    """Resolve an unchanged visible cell to a copyable raw GDS structure."""
    current = library
    current_name = name
    seen: set[tuple[int, str]] = set()
    while True:
        key = (id(current), current_name)
        if key in seen:
            return None
        seen.add(key)

        if isinstance(current, _GdsRawCellSource):
            if current.can_copy_raw_struct(current_name):
                return current, current_name
            return None

        if not isinstance(current, IBorrowing):
            return None
        source_cell = current.source_cell(current_name)
        if source_cell is None:
            return None
        source, source_name = source_cell
        if source_name != current_name:
            return None
        current = source
        current_name = source_name


def _get_write_info(
        library: Mapping[str, Pattern] | ILibraryView,
        *,
        meters_per_unit: float | None,
        logical_units_per_unit: float | None,
        library_name: str | None,
        ) -> tuple[float, float, str]:
    if meters_per_unit is not None and logical_units_per_unit is not None and library_name is not None:
        return meters_per_unit, logical_units_per_unit, library_name

    infos: list[dict[str, Any]] = []
    stack: list[Mapping[str, Pattern] | ILibraryView] = [library]
    seen: set[int] = set()
    while stack:
        current = stack.pop()
        if id(current) in seen:
            continue
        seen.add(id(current))
        if isinstance(current, _GdsInfoSource) and isinstance(current.library_info, dict):
            infos.append(current.library_info)
        if isinstance(current, IBorrowing):
            stack.extend(reversed(current.borrowed_sources()))

    if infos:
        unit_pairs = {(info['meters_per_unit'], info['logical_units_per_unit']) for info in infos}
        if len(unit_pairs) > 1:
            raise LibraryError('Merged lazy GDS sources must have identical units before writing')
        info = infos[0]
        meters = info['meters_per_unit'] if meters_per_unit is None else meters_per_unit
        logical = info['logical_units_per_unit'] if logical_units_per_unit is None else logical_units_per_unit
        name = info['name'] if library_name is None else library_name
        return meters, logical, name

    if meters_per_unit is None:
        raise LibraryError('meters_per_unit is required when writing a library without GDS metadata')
    logical = 1 if logical_units_per_unit is None else logical_units_per_unit
    name = 'masque-klamath' if library_name is None else library_name
    return meters_per_unit, logical, name


def write(
        library: Mapping[str, Pattern] | ILibraryView,
        stream: IO[bytes],
        meters_per_unit: float | None = None,
        logical_units_per_unit: float | None = None,
        library_name: str | None = None,
        ) -> None:
    """Write an eager or source-backed library to a GDSII stream."""
    meters_per_unit, logical_units_per_unit, library_name = _get_write_info(
        library,
        meters_per_unit=meters_per_unit,
        logical_units_per_unit=logical_units_per_unit,
        library_name=library_name,
    )

    klamath._write_header(stream, meters_per_unit, logical_units_per_unit, library_name)

    names = library.source_order() if isinstance(library, ILibraryView) else tuple(library)
    for name in names:
        if isinstance(library, ILibraryView):
            raw_struct = _resolve_raw_struct(library, name)
            if raw_struct is not None:
                raw_source, source_name = raw_struct
                stream.write(raw_source.raw_struct_bytes(source_name))
                continue

        if isinstance(library, IMaterializable):
            pat = library.materialize(name, persist=False)
        else:
            pat = library[name]
        klamath._write_pattern_struct(stream, name, pat)

    klamath._write_footer(stream)


def writefile(
        library: Mapping[str, Pattern] | ILibraryView,
        filename: str | pathlib.Path,
        meters_per_unit: float | None = None,
        logical_units_per_unit: float | None = None,
        library_name: str | None = None,
        ) -> None:
    """Write an eager or source-backed library to a path, compressing `.gz` files."""
    path = pathlib.Path(filename)
    with tmpfile(path) as base_stream:
        if path.suffix == '.gz':
            stream = cast('IO[bytes]', gzip.GzipFile(
                filename='',
                mtime=0,
                fileobj=base_stream,
                mode='wb',
                compresslevel=6,
                ))
            try:
                write(library, stream, meters_per_unit, logical_units_per_unit, library_name)
            finally:
                stream.close()
        else:
            write(library, base_stream, meters_per_unit, logical_units_per_unit, library_name)