Moved master working set construction into classmethods of WorkingSet.

1 files changed, 44 insertions, 27 deletions

diff --git a/pkg_resources.py b/pkg_resources.py
index 0b987e8f..0120ab4f 100644
--- a/pkg_resources.py
+++ b/pkg_resources.py
@@ -427,6 +427,48 @@ class WorkingSet(object):
         for entry in entries:
             self.add_entry(entry)
 
+    @classmethod
+    def _build_master(cls):
+        """
+        Prepare the master working set.
+        """
+        ws = cls()
+        try:
+            from __main__ import __requires__
+        except ImportError:
+            # The main program does not list any requirements
+            return ws
+
+        # ensure the requirements are met
+        try:
+            ws.require(__requires__)
+        except VersionConflict:
+            return cls._build_from_requirements(__requires__)
+
+        return ws
+
+    @classmethod
+    def _build_from_requirements(cls, req_spec):
+        """
+        Build a working set from a requirement spec. Rewrites sys.path.
+        """
+        # try it without defaults already on sys.path
+        # by starting with an empty path
+        ws = cls([])
+        reqs = parse_requirements(req_spec)
+        dists = working_set.resolve(reqs, Environment())
+        for dist in dists:
+            ws.add(dist)
+
+        # add any missing entries from sys.path
+        for entry in sys.path:
+            if entry not in ws.entries:
+                ws.add_entry(entry)
+
+        # then copy back to sys.path
+        sys.path[:] = ws.entries
+        return ws
+
     def add_entry(self, entry):
         """Add a path item to ``.entries``, finding any distributions on it
 
@@ -2704,33 +2746,8 @@ def _initialize(g):
 _initialize(globals())
 
 # Prepare the master working set and make the ``require()`` API available
-_declare_state('object', working_set=WorkingSet())
-try:
-    # Does the main program list any requirements?
-    from __main__ import __requires__
-except ImportError:
-    # No: just use the default working set based on sys.path
-    pass
-else:
-    # Yes: ensure the requirements are met, by prefixing sys.path if necessary
-    try:
-        working_set.require(__requires__)
-    except VersionConflict:
-        # try it without defaults already on sys.path
-        # by starting with an empty path
-        working_set = WorkingSet([])
-        reqs = parse_requirements(__requires__)
-        dists = working_set.resolve(reqs, Environment())
-        for dist in dists:
-            working_set.add(dist)
-
-        # add any missing entries from sys.path
-        for entry in sys.path:
-            if entry not in working_set.entries:
-                working_set.add_entry(entry)
-
-        # then copy back to sys.path
-        sys.path[:] = working_set.entries
+working_set = WorkingSet._build_master()
+_declare_state('object', working_set=working_set)
 
 require = working_set.require
 iter_entry_points = working_set.iter_entry_points