From f587fe5341a929797f81353d2f241d38d1fc622b Mon Sep 17 00:00:00 2001
From: Jan Petykiewicz <jan@mpxd.net>
Date: Sat, 5 Oct 2024 15:51:58 -0700
Subject: [PATCH] add the toolctx() context manager to simplify temporary
 retool() calls

---
 examples/tutorial/pather.py |  6 ++++++
 masque/builder/pather.py    | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/examples/tutorial/pather.py b/examples/tutorial/pather.py
index a7aad68..101fbb5 100644
--- a/examples/tutorial/pather.py
+++ b/examples/tutorial/pather.py
@@ -265,6 +265,12 @@ def main() -> None:
     # when using pather.retool().
     pather.path_to('VCC', None, -50_000, out_ptype='m1wire')
 
+    # Now extend GND out to x=-50_000, using M2 for a portion of the path.
+    # We can use `pather.toolctx()` to temporarily retool, instead of calling `retool()` twice.
+    with pather.toolctx(M2_tool, keys=['GND']):
+        pather.path_to('GND', None, -40_000)
+    pather.path_to('GND', None, -50_000)
+
     # Save the pather's pattern into our library
     library['Pather_and_BasicTool'] = pather.pattern
 
diff --git a/masque/builder/pather.py b/masque/builder/pather.py
index 17d73b3..cf5aaed 100644
--- a/masque/builder/pather.py
+++ b/masque/builder/pather.py
@@ -2,9 +2,10 @@
 Manual wire/waveguide routing (`Pather`)
 """
 from typing import Self
-from collections.abc import Sequence, MutableMapping, Mapping
+from collections.abc import Sequence, MutableMapping, Mapping, Iterator
 import copy
 import logging
+from contextlib import contextmanager
 from pprint import pformat
 
 import numpy
@@ -281,6 +282,37 @@ class Pather(Builder):
                 self.tools[key] = tool
         return self
 
+    @contextmanager
+    def toolctx(
+            self,
+            tool: Tool,
+            keys: str | Sequence[str | None] | None = None,
+            ) -> Iterator[Self]:
+        """
+          Context manager for temporarily `retool`-ing and reverting the `retool`
+        upon exiting the context.
+
+        Args:
+            tool: The new `Tool` to use for the given ports.
+            keys: Which ports the tool should apply to. `None` indicates the default tool,
+                used when there is no matching entry in `self.tools` for the port in question.
+
+        Returns:
+            self
+        """
+        if keys is None or isinstance(keys, str):
+            keys = [keys]
+        saved_tools = {kk: self.tools.get(kk, None) for kk in keys}      # If not in self.tools, save `None`
+        try:
+            yield self.retool(tool=tool, keys=keys)
+        finally:
+            for kk, tt in saved_tools.items():
+                if tt is None:
+                    # delete if present
+                    self.tools.pop(tt, None)
+                else:
+                    self.tools[kk] = tt
+
     def path(
             self,
             portspec: str,