masque/masque/file/gdsii/lazy_write.py

"""
GDS write helpers for source-backed lazy GDS views.

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

import klamath

from . import klamath as gdsii
from ..utils import tmpfile
from ...error import LibraryError
from ...library import IBorrowing, ILibraryView
from ...library.overlay import _materialize_detached_pattern

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 source_order(self) -> tuple[str, ...]: ...

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

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


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 or logical_units_per_unit is None or library_name is None:
        raise LibraryError('meters_per_unit, logical_units_per_unit, and library_name are required for non-GDS-backed lazy writes')
    return meters_per_unit, logical_units_per_unit, library_name


def _write_pattern_struct(stream: IO[bytes], name: str, pat: Pattern) -> None:
    elements: list[klamath.elements.Element] = []
    elements += gdsii._shapes_to_elements(pat.shapes)
    elements += gdsii._labels_to_texts(pat.labels)
    elements += gdsii._mrefs_to_grefs(pat.refs)
    klamath.library.write_struct(stream, name=name.encode('ASCII'), elements=elements)


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:
    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,
    )

    header = klamath.library.FileHeader(
        name=library_name.encode('ASCII'),
        user_units_per_db_unit=logical_units_per_unit,
        meters_per_db_unit=meters_per_unit,
    )
    header.write(stream)

    if isinstance(library, _GdsRawCellSource):
        for name in library.source_order():
            if library.can_copy_raw_struct(name):
                stream.write(library.raw_struct_bytes(name))
            else:
                _write_pattern_struct(stream, name, _materialize_detached_pattern(cast('ILibraryView', library), name))
        klamath.records.ENDLIB.write(stream, None)
        return

    gdsii.write(cast('Mapping[str, Pattern]', library), stream, meters_per_unit, logical_units_per_unit, library_name)


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:
    path = pathlib.Path(filename)

    with tmpfile(path) as base_stream:
        streams: tuple[Any, ...] = (base_stream,)
        if path.suffix == '.gz':
            stream = cast('IO[bytes]', gzip.GzipFile(filename='', mtime=0, fileobj=base_stream, mode='wb', compresslevel=6))
            streams = (stream,) + streams
        else:
            stream = base_stream

        try:
            write(
                library,
                stream,
                meters_per_unit=meters_per_unit,
                logical_units_per_unit=logical_units_per_unit,
                library_name=library_name,
            )
        finally:
            for ss in streams:
                ss.close()