From 39b008df26a449c737caa41238725eb4ed47fd04 Mon Sep 17 00:00:00 2001
From: yardasol <yardasol@grinnell.edu>
Date: Tue, 12 Jul 2022 16:07:59 -0500
Subject: [PATCH] Remove __enter__, __exit__ methods from TransportOperator
 This commit removes the context manager for switching directories from the
 TransportOperator class, and moves this functionality to a module-level
 function, change_directory.

---
 openmc/deplete/abc.py | 45 +++++++++++++++++++++++++++++--------------
 1 file changed, 31 insertions(+), 14 deletions(-)

diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py
index a4f80e3d6..a7cde4afa 100644
--- a/openmc/deplete/abc.py
+++ b/openmc/deplete/abc.py
@@ -10,6 +10,7 @@ from collections.abc import Iterable, Callable
 from copy import deepcopy
 from inspect import signature
 from numbers import Real, Integral
+from contextlib import contextmanager
 import os
 from pathlib import Path
 import sys
@@ -56,6 +57,23 @@ except AttributeError:
     # Can't set __doc__ on properties on Python 3.4
     pass
 
+@contextmanager
+def change_directory(output_dir):
+    """
+    Helper function for managing the current directory.
+
+    Parameters
+    ----------
+    output_dir : pathlib.Path
+        Directory to switch to.
+    """
+    _orig_dir  = os.getcwd()
+    try:
+        output_dir.mkdir(exist_ok=True)
+        os.chdir(output_dir)
+        yield
+    finally:
+        os.chdir(_orig_dir)
 
 class TransportOperator(ABC):
     """Abstract class defining a transport operator
@@ -87,9 +105,14 @@ class TransportOperator(ABC):
         Initial atom density [atoms/cm^3] to add for nuclides that are zero
         in initial condition to ensure they exist in the decay chain.
         Only done for nuclides with reaction rates.
+    output_dir : pathlib.Path
+        Path to output directory to save results.
     prev_res : Results or None
         Results from a previous depletion calculation. ``None`` if no
         results are to be used.
+    chain : openmc.deplete.Chain
+        The depletion chain information necessary to form matrices and tallies.
+
     """
     def __init__(self, chain_file, fission_q=None, dilute_initial=1.0e3,
                  prev_results=None):
@@ -133,18 +156,6 @@ class TransportOperator(ABC):
 
         """
 
-    def __enter__(self):
-        # Save current directory and move to specific output directory
-        self._orig_dir = os.getcwd()
-        self.output_dir.mkdir(exist_ok=True)
-        os.chdir(self.output_dir)
-
-        return self.initial_condition()
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.finalize()
-        os.chdir(self._orig_dir)
-
     @property
     def output_dir(self):
         return self._output_dir
@@ -775,7 +786,8 @@ class Integrator(ABC):
 
             .. versionadded:: 0.13.1
         """
-        with self.operator as conc:
+        with change_directory(self.operator.output_dir):
+            conc = self.operator.initial_condition()
             t, self._i_res = self._get_start_data()
 
             for i, (dt, source_rate) in enumerate(self):
@@ -814,6 +826,8 @@ class Integrator(ABC):
                          source_rate, self._i_res + len(self), proc_time)
             self.operator.write_bos_data(len(self) + self._i_res)
 
+        self.operator.finalize()
+
 
 @add_params
 class SIIntegrator(Integrator):
@@ -931,7 +945,8 @@ class SIIntegrator(Integrator):
 
             .. versionadded:: 0.13.1
         """
-        with self.operator as conc:
+        with change_directory(self.operator.output_dir):
+            conc = self.operator.initial_condition()
             t, self._i_res = self._get_start_data()
 
             for i, (dt, p) in enumerate(self):
@@ -967,6 +982,8 @@ class SIIntegrator(Integrator):
                          p, self._i_res + len(self), proc_time)
             self.operator.write_bos_data(self._i_res + len(self))
 
+        self.operator.finalize()
+
 
 class DepSystemSolver(ABC):
     r"""Abstract class for solving depletion equations