use lazy-loading in example

examples/check.py

@@ -21,8 +21,8 @@ connectivity = [
 cells, props = oasis.readfile('connectivity.oas')
 topcell = cells['top']
 
-polys, labels = snarled.interfaces.masque.read_cell(topcell, connectivity)
-nets_info = snarled.trace_connectivity_preloaded(polys, labels, connectivity)
+get_layer = snarled.interfaces.masque.prepare_cell(topcell)
+nets_info = snarled.trace_connectivity(get_layer, connectivity)
 
 print('\nFinal nets:')
 print([kk for kk in sorted(nets_info.nets.keys()) if isinstance(kk.name, str)])

snarled/interfaces/masque.py

@@ -1,7 +1,7 @@
 """
 Functionality for extracting geometry and label info from `masque` patterns.
 """
-from typing import Sequence, Dict, List, Any, Tuple, Optional, Mapping
+from typing import Sequence, Dict, List, Any, Tuple, Optional, Mapping, Callable
 from collections import defaultdict
 
 import numpy
@@ -14,6 +14,80 @@ from ..types import layer_t
 from ..utils import connectivity2layers
 
 
+def prepare_cell(
+        cell: Pattern,
+        label_mapping: Optional[Mapping[layer_t, layer_t]] = None,
+        ) -> Callable[[layer_t], Tuple[
+            List[NDArray[numpy.float64]],
+            List[Tuple[float, float, str]]
+            ]]:
+    """
+    Generate a function for extracting `polys` and `labels` from a `masque.Pattern`.
+    The returned function can be passed to `snarled.trace_connectivity`.
+
+    Args:
+        cell: A `masque` `Pattern` object. Usually your topcell.
+        label_mapping: A mapping of `{label_layer: metal_layer}`. This allows labels
+            to refer to nets on metal layers without the labels themselves being on
+            that layer.
+            Default `None` reads labels from the same layer as the geometry.
+
+    Returns:
+        `get_layer` function, to be passed to `snarled.trace_connectivity`.
+    """
+
+    def get_layer(
+            layer: layer_t,
+            ) -> Tuple[
+                List[NDArray[numpy.float64]],
+                List[Tuple[float, float, str]]
+                ]:
+
+        if label_mapping is None:
+            label_layers = {layer: layer}
+        else:
+            label_layers = {label_layer for label_layer, metal_layer in label_mapping.items()
+                            if metal_layer == layer}
+
+        subset = cell.deepcopy().subset(          # TODO add single-op subset-and-copy, to avoid copying unwanted stuff
+            shapes_func=lambda ss: ss.layer == layer,
+            labels_func=lambda ll: ll.layer in label_layers,
+            subpatterns_func=lambda ss: True,
+            )
+
+        flat = subset.flatten()
+
+        # load polygons
+        polys = []
+        for ss in flat.shapes:
+            assert(isinstance(ss, Polygon))
+
+            if ss.repetition is None:
+                displacements = [(0, 0)]
+            else:
+                displacements = ss.repetition.displacements
+
+            for displacement in displacements:
+                polys.append(
+                    ss.vertices + ss.offset + displacement
+                    )
+
+        # load metal labels
+        labels = []
+        for ll in flat.labels:
+            if ll.repetition is None:
+                displacements = [(0, 0)]
+            else:
+                displacements = ll.repetition.displacements
+
+            for displacement in displacements:
+                offset = ll.offset + displacement
+                labels.append((*offset, ll.string))
+
+        return polys, labels
+    return get_layer
+
+
 def read_cell(
         cell: Pattern,
         connectivity: Sequence[Tuple[layer_t, Optional[layer_t], layer_t]],
@@ -41,57 +115,11 @@ def read_cell(
     metal_layers, via_layers = connectivity2layers(connectivity)
     poly_layers = metal_layers | via_layers
 
-    if label_mapping is None:
-        label_mapping = {layer: layer for layer in metal_layers}
-    label_layers = {label_layer for label_layer in label_mapping.keys()}
+    get_layer = prepare_cell(cell, label_mapping)
 
-    cell = cell.deepcopy().subset(
-        shapes_func=lambda ss: ss.layer in poly_layers,
-        labels_func=lambda ll: ll.layer in label_layers,
-        subpatterns_func=lambda ss: True,
-        )
-
-    # load polygons
     polys = defaultdict(list)
+    labels = defaultdict(list)
     for layer in poly_layers:
-        shapes_hier = cell.subset(
-            shapes_func=lambda ss: ss.layer == layer,
-            subpatterns_func=lambda ss: True,
-            )
-        shapes = shapes_hier.flatten().shapes
+        polys[layer], labels[layer] = get_layer(layer)
 
-        for ss in shapes:
-            assert(isinstance(ss, Polygon))
-
-            if ss.repetition is None:
-                displacements = [(0, 0)]
-            else:
-                displacements = ss.repetition.displacements
-
-            for displacement in displacements:
-                polys[ss.layer].append(
-                    ss.vertices + ss.offset + displacement
-                    )
-
-    # load metal labels
-    metal_labels = defaultdict(list)
-    for label_layer, metal_layer in label_mapping.items():
-        labels_hier = cell.subset(
-            labels_func=lambda ll: ll.layer == label_layer,
-            subpatterns_func=lambda ss: True,
-            )
-        labels = labels_hier.flatten().labels
-
-        for ll in labels:
-            if ll.repetition is None:
-                displacements = [(0, 0)]
-            else:
-                displacements = ll.repetition.displacements
-
-            for displacement in displacements:
-                offset = ll.offset + displacement
-                metal_labels[metal_layer].append(
-                    (*offset, ll.string)
-                    )
-
-    return polys, metal_labels
+    return polys, labels