masque/masque/test/test_file_roundtrip.py

from pathlib import Path
from typing import cast
import pytest
from numpy.testing import assert_allclose

from ..pattern import Pattern
from ..library import Library
from ..shapes import Path as MPath, Circle, Polygon
from ..repetition import Grid, Arbitrary

def create_test_library(for_gds: bool = False) -> Library:
    from ..file import gdsii
    if not for_gds:
        from ..file import oasis
    lib = Library()

    # 1. Polygons
    pat_poly = Pattern()
    pat_poly.polygon((1, 0), vertices=[[0, 0], [10, 0], [5, 10]])
    lib["polygons"] = pat_poly

    # 2. Paths with different endcaps
    pat_paths = Pattern()
    # Flush
    pat_paths.path((2, 0), vertices=[[0, 0], [20, 0]], width=2, cap=MPath.Cap.Flush)
    # Square
    pat_paths.path((2, 1), vertices=[[0, 10], [20, 10]], width=2, cap=MPath.Cap.Square)
    # Circle (Only for GDS)
    if for_gds:
        pat_paths.path((2, 2), vertices=[[0, 20], [20, 20]], width=2, cap=MPath.Cap.Circle)
    # SquareCustom
    pat_paths.path((2, 3), vertices=[[0, 30], [20, 30]], width=2, cap=MPath.Cap.SquareCustom, cap_extensions=(1, 5))
    lib["paths"] = pat_paths

    # 3. Circles (only for OASIS or polygonized for GDS)
    pat_circles = Pattern()
    if for_gds:
        # GDS writer calls to_polygons() for non-supported shapes,
        # but we can also pre-polygonize
        pat_circles.shapes[(3, 0)].append(Circle(radius=5, offset=(10, 10)).to_polygons()[0])
    else:
        pat_circles.shapes[(3, 0)].append(Circle(radius=5, offset=(10, 10)))
    lib["circles"] = pat_circles

    # 4. Refs with repetitions
    pat_refs = Pattern()
    # Simple Ref
    pat_refs.ref("polygons", offset=(0, 0))
    # Ref with Grid repetition
    pat_refs.ref("polygons", offset=(100, 0), repetition=Grid(a_vector=(20, 0), a_count=3, b_vector=(0, 20), b_count=2))
    # Ref with Arbitrary repetition
    pat_refs.ref("polygons", offset=(0, 100), repetition=Arbitrary(displacements=[[0, 0], [10, 20], [30, -10]]))
    lib["refs"] = pat_refs

    # 5. Shapes with repetitions (OASIS only, must be wrapped for GDS)
    pat_rep_shapes = Pattern()
    poly_rep = Polygon(vertices=[[0, 0], [5, 0], [5, 5], [0, 5]], repetition=Grid(a_vector=(10, 0), a_count=5))
    pat_rep_shapes.shapes[(4, 0)].append(poly_rep)
    lib["rep_shapes"] = pat_rep_shapes

    if for_gds:
        lib.wrap_repeated_shapes()

    return lib

def test_gdsii_full_roundtrip(tmp_path: Path) -> None:
    from ..file import gdsii
    lib = create_test_library(for_gds=True)
    gds_file = tmp_path / "full_test.gds"
    gdsii.writefile(lib, gds_file, meters_per_unit=1e-9)

    read_lib, _ = gdsii.readfile(gds_file)

    # Check existence
    for name in lib:
        assert name in read_lib

    # Check Paths
    read_paths = read_lib["paths"]
    # Check caps (GDS stores them as path_type)
    # Order might be different depending on how they were written,
    # but here they should match the order they were added if dict order is preserved.
    # Actually, they are grouped by layer.
    p_flush = cast("MPath", read_paths.shapes[(2, 0)][0])
    assert p_flush.cap == MPath.Cap.Flush

    p_square = cast("MPath", read_paths.shapes[(2, 1)][0])
    assert p_square.cap == MPath.Cap.Square

    p_circle = cast("MPath", read_paths.shapes[(2, 2)][0])
    assert p_circle.cap == MPath.Cap.Circle

    p_custom = cast("MPath", read_paths.shapes[(2, 3)][0])
    assert p_custom.cap == MPath.Cap.SquareCustom
    assert p_custom.cap_extensions is not None
    assert_allclose(p_custom.cap_extensions, (1, 5))

    # Check Refs with repetitions
    read_refs = read_lib["refs"]
    assert len(read_refs.refs["polygons"]) >= 3 # Simple, Grid (becomes 1 AREF), Arbitrary (becomes 3 SREFs)

    # AREF check
    arefs = [r for r in read_refs.refs["polygons"] if r.repetition is not None]
    assert len(arefs) == 1
    assert isinstance(arefs[0].repetition, Grid)
    assert arefs[0].repetition.a_count == 3
    assert arefs[0].repetition.b_count == 2

    # Check wrapped shapes
    # lib.wrap_repeated_shapes() created new patterns
    # Original pattern "rep_shapes" now should have a Ref
    assert len(read_lib["rep_shapes"].refs) > 0

def test_oasis_full_roundtrip(tmp_path: Path) -> None:
    pytest.importorskip("fatamorgana")
    from ..file import oasis
    lib = create_test_library(for_gds=False)
    oas_file = tmp_path / "full_test.oas"
    oasis.writefile(lib, oas_file, units_per_micron=1000)

    read_lib, _ = oasis.readfile(oas_file)

    # Check existence
    for name in lib:
        assert name in read_lib

    # Check Circle
    read_circles = read_lib["circles"]
    assert isinstance(read_circles.shapes[(3, 0)][0], Circle)
    assert read_circles.shapes[(3, 0)][0].radius == 5

    # Check Path caps
    read_paths = read_lib["paths"]
    assert cast("MPath", read_paths.shapes[(2, 0)][0]).cap == MPath.Cap.Flush
    assert cast("MPath", read_paths.shapes[(2, 1)][0]).cap == MPath.Cap.Square
    # OASIS HalfWidth is Square. masque's Square is also HalfWidth extension.
    # Wait, Circle cap in OASIS?
    # masque/file/oasis.py:
    # path_cap_map = {
    #     PathExtensionScheme.Flush: Path.Cap.Flush,
    #     PathExtensionScheme.HalfWidth: Path.Cap.Square,
    #     PathExtensionScheme.Arbitrary: Path.Cap.SquareCustom,
    #     }
    # It seems Circle cap is NOT supported in OASIS by masque currently.
    # Let's verify what happens with Circle cap in OASIS write.
    # _shapes_to_elements in oasis.py:
    # path_type = next(k for k, v in path_cap_map.items() if v == shape.cap)
    # This will raise StopIteration if Circle is not in path_cap_map.

    # Check Shape repetition
    read_rep_shapes = read_lib["rep_shapes"]
    poly = read_rep_shapes.shapes[(4, 0)][0]
    assert poly.repetition is not None
    assert isinstance(poly.repetition, Grid)
    assert poly.repetition.a_count == 5