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, cast
import gzip
import logging
import pathlib

import klamath

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

if TYPE_CHECKING:
    from collections.abc import Iterator, Mapping

    from ...pattern import Pattern


logger = logging.getLogger(__name__)


def _iter_library_infos(library: Mapping[str, Pattern] | ILibraryView) -> Iterator[dict[str, Any]]:
    info = getattr(library, 'library_info', None)
    if isinstance(info, dict):
        yield info
    if isinstance(library, OverlayLibrary):
        for layer in library._layers:
            yield from _iter_library_infos(layer.library)


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(_iter_library_infos(library))
    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 _can_copy_raw_cell(library: Mapping[str, Pattern] | ILibraryView, name: str) -> bool:
    can_copy = getattr(library, 'can_copy_raw_struct', None)
    if not callable(can_copy):
        return False
    return bool(can_copy(name))


def _raw_struct_bytes(library: Mapping[str, Pattern] | ILibraryView, name: str) -> bytes:
    reader = getattr(library, 'raw_struct_bytes', None)
    if not callable(reader):
        raise TypeError('raw_struct_bytes')
    return cast('bytes', reader(name))


def _can_copy_overlay_cell(library: OverlayLibrary, name: str, entry: _SourceEntry) -> bool:
    layer = library._layers[entry.layer_index]
    if name != entry.source_name:
        return False
    if not _can_copy_raw_cell(layer.library, entry.source_name):
        return False
    children = layer.child_graph.get(entry.source_name, set())
    return all(library._effective_target(layer, child) == child for child in children)


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, OverlayLibrary):
        for name in library.source_order():
            entry = library._entries[name]
            if isinstance(entry, _SourceEntry) and _can_copy_overlay_cell(library, name, entry):
                layer = library._layers[entry.layer_index]
                stream.write(_raw_struct_bytes(layer.library, entry.source_name))
            else:
                _write_pattern_struct(stream, name, library._materialize_pattern(name, persist=False))
        klamath.records.ENDLIB.write(stream, None)
        return

    if hasattr(library, 'raw_struct_bytes'):
        for name in library.source_order():
            if _can_copy_raw_cell(library, name):
                stream.write(_raw_struct_bytes(library, 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()