add --local option to freeze; fixes #9

pip/commands/freeze.py

@@ -5,6 +5,7 @@ import pip
 from pip.req import InstallRequirement
 from pip.log import logger
 from pip.basecommand import Command
+from pip.util import dist_location, is_local
 
 class FreezeCommand(Command):
     name = 'freeze'
@@ -27,10 +28,17 @@ class FreezeCommand(Command):
             default=[],
             metavar='URL',
             help='URL for finding packages, which will be added to the frozen requirements file')
+        self.parser.add_option(
+            '-l', '--local',
+            dest='local',
+            action='store_true',
+            default=False,
+            help='If in a virtualenv, do not report globally-installed packages')
 
     def run(self, options, args):
         requirement = options.requirement
         find_links = options.find_links or []
+        local_only = options.local
         ## FIXME: Obviously this should be settable:
         find_tags = False
         skip_match = None
@@ -54,6 +62,8 @@ class FreezeCommand(Command):
             f.write('-f %s\n' % link)
         installations = {}
         for dist in pkg_resources.working_set:
+            if local_only and not is_local(dist_location(dist)):
+                continue
             if dist.key in ('setuptools', 'pip', 'python'):
                 ## FIXME: also skip virtualenv?
                 continue

pip/req.py

@@ -19,8 +19,8 @@ from pip.log import logger
 from pip.util import display_path, rmtree, format_size
 from pip.util import splitext, ask, backup_dir
 from pip.util import url_to_filename, filename_to_url
-from pip.util import is_url, is_filename
-from pip.util import renames, normalize_path
+from pip.util import is_url, is_filename, is_local
+from pip.util import renames, normalize_path, egg_link_path
 from pip.util import make_path_relative, is_svn_page, file_contents
 from pip.util import has_leading_dir, split_leading_dir
 from pip.util import get_file_content
@@ -378,22 +378,12 @@ execfile(__file__)
             raise UninstallationError("Cannot uninstall requirement %s, not installed" % (self.name,))
         dist = self.satisfied_by or self.conflicts_with
 
-        # we restrict uninstallation to sys.prefix only if sys.real_prefix exists
-        # i.e. in a virtualenv, only uninstall within virtualenv. Otherwise uninstall anything.
-        paths_to_remove = UninstallPathSet(dist, getattr(sys, 'real_prefix', None) and sys.prefix)
+        paths_to_remove = UninstallPathSet(dist)
 
         pip_egg_info_path = os.path.join(dist.location,
                                          dist.egg_name()) + '.egg-info'
         easy_install_egg = dist.egg_name() + '.egg'
-        # This won't find a globally-installed develop egg if
-        # we're in a virtualenv.
-        # (There doesn't seem to be any metadata in the
-        # Distribution object for a develop egg that points back
-        # to its .egg-link and easy-install.pth files).  That's
-        # OK, because we restrict ourselves to making changes
-        # within sys.prefix anyway.
-        develop_egg_link = os.path.join(site_packages,
-                                        dist.project_name) + '.egg-link'
+        develop_egg_link = egg_link_path(dist)
         if os.path.exists(pip_egg_info_path):
             # package installed by pip
             paths_to_remove.add(pip_egg_info_path)
@@ -1373,11 +1363,10 @@ def parse_editable(editable_req, default_vcs=None):
 class UninstallPathSet(object):
     """A set of file paths to be removed in the uninstallation of a
     requirement."""
-    def __init__(self, dist, restrict_to_prefix=None):
+    def __init__(self, dist):
         self.paths = set()
         self._refuse = set()
         self.pth = {}
-        self.prefix = restrict_to_prefix and normalize_path(restrict_to_prefix)
         self.dist = dist
         self.location = normalize_path(dist.location)
         self.save_dir = None
@@ -1389,15 +1378,12 @@ class UninstallPathSet(object):
         False otherwise.
 
         """
-        if self.prefix:
-            return path.startswith(self.prefix)
-        return True
+        return is_local(path)
         
     def _can_uninstall(self):
         if not self._permitted(self.location):
             logger.notify("Not uninstalling %s at %s, outside environment %s"
-                          % (self.dist.project_name, self.location,
-                             self.prefix))
+                          % (self.dist.project_name, self.location, sys.prefix))
             return False
         return True
 

pip/util.py

@@ -8,6 +8,7 @@ import urllib2
 import re
 from pip.backwardcompat import WindowsError
 from pip.exceptions import InstallationError
+from pip.locations import site_packages
 
 __all__ = ['rmtree', 'display_path', 'backup_dir',
            'find_command', 'splitext', 'ask', 'Inf',
@@ -310,3 +311,40 @@ def renames(old, new):
             os.removedirs(head)
         except OSError:
             pass
+
+def is_local(path):
+    """
+    Return True if path is within sys.prefix, if we're running in a virtualenv.
+
+    If we're not in a virtualenv, all paths are considered "local."
+
+    """
+    if not hasattr(sys, 'real_prefix'):
+        return True
+    return normalize_path(path).startswith(normalize_path(sys.prefix))
+
+def egg_link_path(dist):
+    """
+    Return the path where we'd expect to find a .egg-link file for
+    this distribution. (There doesn't seem to be any metadata in the
+    Distribution object for a develop egg that points back to its
+    .egg-link and easy-install.pth files).
+
+    This won't find a globally-installed develop egg if we're in a
+    virtualenv. 
+
+    """
+    return os.path.join(site_packages, dist.project_name) + '.egg-link'
+
+def dist_location(dist):
+    """
+    Get the site-packages location of this distribution. Generally
+    this is dist.location, except in the case of develop-installed
+    packages, where dist.location is the source code location, and we
+    want to know where the egg-link file is.
+
+    """
+    egg_link = egg_link_path(dist)
+    if os.path.exists(egg_link):
+        return egg_link
+    return dist.location

tests/test_freeze.txt

@@ -111,7 +111,7 @@ Now what about Mercurial::
     -e hg+http://bitbucket.org/jezdez/django-dbtemplates/@...#egg=django_dbtemplates-...
     ...
 
-Heck, now look in the Bazaar:
+Heck, now look in the Bazaar::
 
     >>> reset_env()
     >>> env = get_env()
@@ -132,3 +132,22 @@ Heck, now look in the Bazaar:
     -f bzr+http://bazaar.launchpad.net/...django-wikiapp/django-wikiapp/release-0.1/#egg=django-wikiapp
     -e bzr+http://bazaar.launchpad.net/...django-wikiapp/django-wikiapp/release-0.1/@...#egg=django_wikiapp-...
     ...
+
+
+Test that wsgiref (from global site-packages) is reported normally, but not with --local::
+
+    >>> reset_env()
+    >>> result = run_pip('install', 'initools==0.2')
+    >>> result = run_pip('freeze', expect_stderr=True)
+    >>> print result
+    Script result: ...ython... pip.main() .../test-scratch freeze
+    -- stdout: --------------------
+    INITools==0.2
+    wsgiref==...
+    <BLANKLINE>
+    >>> result2 = run_pip('freeze', '--local', expect_stderr=True)
+    >>> print result2
+    Script result: ...ython... pip.main() .../test-scratch freeze --local
+    -- stdout: --------------------
+    INITools==0.2
+    <BLANKLINE>