Implemented a context manager for os.chdir

src/mepo/command/update-state.py

@@ -1,7 +1,6 @@
-import os
-
 from ..state import MepoState
 from ..component import MepoComponent
+from ..utilities.chdir import chdir as mepo_chdir
 
 def run(args):
     '''Permanently convert mepo1 state to mepo2 state'''
@@ -18,10 +17,8 @@ def run(args):
         # This needs to be run from the fixture directory so that
         # MepoComponent::__set_original_version
         # picks the right repo for setting version
-        cwd = os.getcwd()
-        os.chdir(MepoState.get_root_dir())
-        comp_new = MepoComponent().to_component(tmp_name, tmp_details, None)
-        os.chdir(cwd)
+        with mepo_chdir(MepoState.get_root_dir()):
+            comp_new = MepoComponent().to_component(tmp_name, tmp_details, None)
         allcomps.append(comp_new)
     # Write new state
     MepoState.write_state(allcomps)

src/mepo/state.py

@@ -12,6 +12,7 @@
 from .utilities import colors
 from .utilities.exceptions import StateDoesNotExistError
 from .utilities.exceptions import StateAlreadyInitializedError
+from .utilities.chdir import chdir as mepo_chdir
 
 class MepoState(object):
 
@@ -134,7 +135,5 @@ def write_state(cls, state_details):
         if os.path.isfile(state_fileptr_fullpath):
             os.remove(state_fileptr_fullpath)
         #os.symlink(new_state_file, state_fileptr_fullpath)
-        curr_dir=os.getcwd()
-        os.chdir(state_dir)
-        os.symlink(state_file_name, state_fileptr)
-        os.chdir(curr_dir)
+        with mepo_chdir(state_dir):
+            os.symlink(state_file_name, state_fileptr)

src/mepo/utilities/chdir.py

@@ -0,0 +1,11 @@
+import os
+from contextlib import contextmanager
+
+@contextmanager
+def chdir(path):
+    cwd = os.getcwd()
+    os.chdir(path)
+    try:
+        yield
+    finally:
+        os.chdir(cwd)