Extract _patch_usage and re-implement as a context manager.

setuptools/command/easy_install.py

@@ -34,6 +34,7 @@ import textwrap
 import warnings
 import site
 import struct
+import contextlib
 
 from setuptools import Command
 from setuptools.sandbox import run_setup
@@ -2119,39 +2120,42 @@ def bootstrap():
 def main(argv=None, **kw):
     from setuptools import setup
     from setuptools.dist import Distribution
-    import distutils.core
-
-    USAGE = textwrap.dedent("""
-        usage: %(script)s [options] requirement_or_url ...
-           or: %(script)s --help
-        """).lstrip()
-
-    def gen_usage(script_name):
-        return USAGE % dict(
-            script=os.path.basename(script_name),
-        )
-
-    def with_ei_usage(f):
-        old_gen_usage = distutils.core.gen_usage
-        try:
-            distutils.core.gen_usage = gen_usage
-            return f()
-        finally:
-            distutils.core.gen_usage = old_gen_usage
 
     class DistributionWithoutHelpCommands(Distribution):
         common_usage = ""
 
         def _show_help(self, *args, **kw):
-            with_ei_usage(lambda: Distribution._show_help(self, *args, **kw))
+            with _patch_usage():
+                Distribution._show_help(self, *args, **kw)
 
     if argv is None:
         argv = sys.argv[1:]
 
-    with_ei_usage(
-        lambda: setup(
+    with _patch_usage():
+        setup(
             script_args=['-q', 'easy_install', '-v'] + argv,
             script_name=sys.argv[0] or 'easy_install',
             distclass=DistributionWithoutHelpCommands, **kw
         )
-    )
+
+
+@contextlib.contextmanager
+def _patch_usage():
+    import distutils.core
+    USAGE = textwrap.dedent("""
+        usage: %(script)s [options] requirement_or_url ...
+           or: %(script)s --help
+        """).lstrip()
+
+    def gen_usage(script_name):
+        return USAGE % dict(
+            script=os.path.basename(script_name),
+        )
+
+    saved = distutils.core.gen_usage
+    distutils.core.gen_usage = gen_usage
+    try:
+        yield
+    finally:
+        distutils.core.gen_usage = saved
+