Provisional monkeypatch fix for windows vs multiprocessing w/nosetests script

1 files changed, 28 insertions, 0 deletions

diff --git a/setup.py b/setup.py
index b77080c..f95858f 100644
--- a/setup.py
+++ b/setup.py
@@ -19,6 +19,11 @@ try:
         },
         test_suite = 'nose.collector',
         )
+    # This is required by multiprocess plugin; on Windows, if
+    # the launch script is not import-safe, spawned processes
+    # will re-run it, resulting in an infinite loop.
+    if sys.platform == 'win32':
+        patch_setuptools()
 except ImportError:
     from distutils.core import setup
     addl_args = dict(
@@ -74,3 +79,26 @@ setup(
         ],
     **addl_args
     )
+
+def patch_setuptools():
+    import re
+    from setuptools.command.easy_install import easy_install
+
+    def wrap_write_script(self, script_name, contents, **kwarg):
+        bad_text = re.compile(
+            "\n"
+            "sys.exit\(\n"
+            "   load_entry_point\(([^\)]+)\)\(\)\n"
+            "\)\n")
+        good_text = (
+            "\n"
+            "if __name__ == '__main__':\n"
+            "    sys.exit(\n"
+            r"        load_entry_point(\1)()\n"
+            "    )\n"
+            )
+        contents = bad_text.sub(good_text, cotnents)
+        return self._write_script(script_name, contents, **kwarg)
+    easy_install._write_script = easy_install.write_script
+    easy_install.write_script = wrap_write_script
+