From 915654bedeee45790a9c251317f113870a81d552 Mon Sep 17 00:00:00 2001
From: Tim Brooks <41971846+timryanb@users.noreply.github.com>
Date: Tue, 20 Feb 2024 11:04:10 -0500
Subject: [PATCH] Updating run_directory API to use more general wraps
 decorator

---
 mphys/scenario.py | 35 +++++++++++++++++++++++------------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/mphys/scenario.py b/mphys/scenario.py
index cb6899da..5015ab78 100644
--- a/mphys/scenario.py
+++ b/mphys/scenario.py
@@ -1,6 +1,17 @@
+from functools import wraps
+
 from .mphys_group import MphysGroup
 from .utils.directory_utils import cd
 
+# Define decorator functions for methods where run directory must be switched before calling
+def switch_run_directory(method):
+    @wraps(method)
+    def wrapped_method(self, *args, **kwargs):
+        with cd(self.options['run_directory']):
+            return method(self, *args, **kwargs)
+
+    return wrapped_method
+
 class Scenario(MphysGroup):
     """
     A group to represent a specific analysis condition or point of the MPhys
@@ -60,21 +71,21 @@ def mphys_add_post_subsystem(self, name, subsystem,
         # setup() to add the builder subsystems before adding these
         self._post_subsystems.append((name, subsystem, promotes_inputs, promotes_outputs, promotes))
 
-    def _solve_nonlinear(self):
-        with cd(self.options['run_directory']):
-            return super()._solve_nonlinear()
+    @switch_run_directory
+    def _solve_nonlinear(self, *args, **kwargs):
+        return super()._solve_nonlinear(*args, **kwargs)
 
-    def _solve_linear(self, mode, rel_systems, scope_out=..., scope_in=...):
-        with cd(self.options['run_directory']):
-            return super()._solve_linear(mode, rel_systems, scope_out, scope_in)
+    @switch_run_directory
+    def _solve_linear(self, *args, **kwargs):
+        return super()._solve_linear(*args, **kwargs)
 
-    def _apply_nonlinear(self):
-        with cd(self.options['run_directory']):
-            return super()._apply_nonlinear()
+    @switch_run_directory
+    def _apply_nonlinear(self, *args, **kwargs):
+        return super()._apply_nonlinear(*args, **kwargs)
 
-    def _apply_linear(self, jac, rel_systems, mode, scope_out=None, scope_in=None):
-        with cd(self.options['run_directory']):
-            return super()._apply_linear(jac, rel_systems, mode, scope_out, scope_in)
+    @switch_run_directory
+    def _apply_linear(self, *args, **kwargs):
+        return super()._apply_linear(*args, **kwargs)
 
     def _mphys_add_pre_coupling_subsystem_from_builder(self, name, builder, scenario_name=None):
         """