TST,STY: Use a context manager for F2PY functions

2 files changed, 25 insertions, 22 deletions

diff --git a/numpy/f2py/tests/test_compile_function.py b/numpy/f2py/tests/test_compile_function.py
index f76fd6448..33e320731 100644
--- a/numpy/f2py/tests/test_compile_function.py
+++ b/numpy/f2py/tests/test_compile_function.py
@@ -54,33 +54,30 @@ def test_f2py_init_compile(extra_args):
         # util.py, but don't actually use build_module() because it has
         # its own invocation of subprocess that circumvents the
         # f2py.compile code block under test
-        try:
-            os.chdir(moddir)
+        with util.switchdir(moddir):
             ret_val = numpy.f2py.compile(
                 fsource,
                 modulename=modname,
                 extra_args=extra_args,
                 source_fn=source_fn
-                )
-        finally:
-            os.chdir(cwd)
+            )
 
-        # check for compile success return value
-        assert_equal(ret_val, 0)
+            # check for compile success return value
+            assert_equal(ret_val, 0)
 
         # we are not currently able to import the Python-Fortran
         # interface module on Windows / Appveyor, even though we do get
         # successful compilation on that platform with Python 3.x
-        if sys.platform != 'win32':
+    if sys.platform != 'win32':
             # check for sensible result of Fortran function; that means
             # we can import the module name in Python and retrieve the
             # result of the sum operation
-            return_check = import_module(modname)
-            calc_result = return_check.foo()
-            assert_equal(calc_result, 15)
+        return_check = import_module(modname)
+        calc_result = return_check.foo()
+        assert_equal(calc_result, 15)
             # Removal from sys.modules, is not as such necessary. Even with
             # removal, the module (dict) stays alive.
-            del sys.modules[modname]
+        del sys.modules[modname]
 
 
 def test_f2py_init_compile_failure():
@@ -113,13 +110,9 @@ def test_f2py_init_compile_bad_cmd():
          b'program test_f2py\nend program test_f2py',])
 def test_compile_from_strings(tmpdir, fsource):
     # Make sure we can compile str and bytes gh-12796
-    cwd = os.getcwd()
-    try:
-        os.chdir(str(tmpdir))
+    with util.switchdir(tmpdir):
         ret_val = numpy.f2py.compile(
-                fsource,
-                modulename='test_compile_from_strings',
-                extension='.f90')
+            fsource,
+            modulename='test_compile_from_strings',
+            extension='.f90')
         assert_equal(ret_val, 0)
-    finally:
-        os.chdir(cwd)
diff --git a/numpy/f2py/tests/util.py b/numpy/f2py/tests/util.py
index c1dc4a536..cfa75c16a 100644
--- a/numpy/f2py/tests/util.py
+++ b/numpy/f2py/tests/util.py
@@ -15,6 +15,7 @@ import atexit
 import textwrap
 import re
 import pytest
+import contextlib
 
 from pathlib import Path
 from numpy.compat import asbytes, asstr
@@ -95,8 +96,8 @@ def build_module(source_files, options=[], skip=[], only=[], module_name=None):
 
     """
 
-    code = ("import sys; sys.path = %s; import numpy.f2py as f2py2e; "
-            "f2py2e.main()" % repr(sys.path))
+    code = ("import sys; sys.path = %s; import numpy.f2py; "
+            "numpy.f2py.main()" % repr(sys.path))
 
     d = get_module_dir()
 
@@ -365,3 +366,12 @@ def getpath(*a):
     # Package root
     d = Path(__file__).parent.parent.resolve()
     return d.joinpath(*a)
+
+@contextlib.contextmanager
+def switchdir(path):
+    curpath = Path.cwd()
+    os.chdir(path)
+    try:
+        yield
+    finally:
+        os.chdir(curpath)