Create custom get_distribution function

Author: deveshks

src/pip/_internal/commands/search.py

@@ -22,7 +22,7 @@
 from pip._internal.network.xmlrpc import PipXmlrpcTransport
 from pip._internal.utils.compat import get_terminal_size
 from pip._internal.utils.logging import indent_log
-from pip._internal.utils.misc import write_output
+from pip._internal.utils.misc import get_distribution, write_output
 
 logger = logging.getLogger(__name__)
 
@@ -127,7 +127,7 @@ def print_results(hits, name_column_width=None, terminal_width=None):
         try:
             write_output(line)
             if name in installed_packages:
-                dist = pkg_resources.get_distribution(name)
+                dist = get_distribution(name)
                 with indent_log():
                     if dist.version == latest:
                         write_output('INSTALLED: %s (latest)', dist.version)

src/pip/_internal/req/req_install.py

@@ -40,6 +40,7 @@
     display_path,
     dist_in_site_packages,
     dist_in_usersite,
+    get_distribution,
     get_installed_version,
     hide_url,
     redact_auth_from_url,
@@ -421,21 +422,23 @@ def check_if_exists(self, use_user_site):
         if self.req is None:
             return
 
-        # Canonicalize requirement name to use normalized
-        # names while searching for already installed packages
-        no_marker = Requirement(str(self.req))
-        no_marker.marker = None
-        no_marker.name = canonicalize_name(no_marker.name)
         # get_distribution() will resolve the entire list of requirements
         # anyway, and we've already determined that we need the requirement
         # in question, so strip the marker so that we don't try to
         # evaluate it.
+        no_marker = Requirement(str(self.req))
+        no_marker.marker = None
+
+        # pkg_resources uses the canonical name to look up packages, but
+        # the name passed passed to get_distribution is not canonicalized
+        # so we have to explicitly convert it to a canonical name
+        no_marker.name = canonicalize_name(no_marker.name)
         try:
             self.satisfied_by = pkg_resources.get_distribution(str(no_marker))
         except pkg_resources.DistributionNotFound:
             return
         except pkg_resources.VersionConflict:
-            existing_dist = pkg_resources.get_distribution(
+            existing_dist = get_distribution(
                 self.req.name
             )
             if use_user_site:
@@ -668,13 +671,11 @@ def uninstall(self, auto_confirm=False, verbose=False):
 
         """
         assert self.req
-        try:
-            dist = pkg_resources.get_distribution(self.req.name)
-        except pkg_resources.DistributionNotFound:
+        dist = get_distribution(self.req.name)
+        if not dist:
             logger.warning("Skipping %s as it is not installed.", self.name)
             return None
-        else:
-            logger.info('Found existing installation: %s', dist)
+        logger.info('Found existing installation: %s', dist)
 
         uninstalled_pathset = UninstallPathSet.from_dist(dist)
         uninstalled_pathset.remove(auto_confirm, verbose)

src/pip/_internal/self_outdated_check.py

@@ -10,7 +10,6 @@
 import os.path
 import sys
 
-from pip._vendor import pkg_resources
 from pip._vendor.packaging import version as packaging_version
 from pip._vendor.six import ensure_binary
 
@@ -25,6 +24,7 @@
 )
 from pip._internal.utils.misc import (
     ensure_dir,
+    get_distribution,
     get_installed_version,
     redact_auth_from_url,
 )
@@ -149,11 +149,10 @@ def was_installed_by_pip(pkg):
     This is used not to display the upgrade message when pip is in fact
     installed by system package manager, such as dnf on Fedora.
     """
-    try:
-        dist = pkg_resources.get_distribution(pkg)
-        return "pip" == get_installer(dist)
-    except pkg_resources.DistributionNotFound:
+    dist = get_distribution(pkg)
+    if not dist:
         return False
+    return "pip" == get_installer(dist)
 
 
 def pip_self_version_check(session, options):

src/pip/_internal/utils/misc.py

@@ -20,6 +20,7 @@
 from pip._vendor import pkg_resources
 # NOTE: retrying is not annotated in typeshed as on 2017-07-17, which is
 #       why we ignore the type on this import.
+from pip._vendor.packaging.utils import canonicalize_name
 from pip._vendor.retrying import retry  # type: ignore
 from pip._vendor.six import PY2, text_type
 from pip._vendor.six.moves import input, map, zip_longest
@@ -480,6 +481,40 @@ def user_test(d):
             ]
 
 
+def search_distribution(req_name):
+
+    # Canonicalize the name before searching in the list of
+    # installed distributions and also while creating the package
+    # dictionary to get the Distribution object
+    req_name = canonicalize_name(req_name)
+    packages = get_installed_distributions(skip=())
+    pkg_dict = {canonicalize_name(p.key): p for p in packages}
+    return pkg_dict.get(req_name)
+
+
+def get_distribution(req_name):
+    """Given a requirement name, return the installed Distribution object"""
+
+    # Search the distribution by looking through the working set
+    dist = search_distribution(req_name)
+
+    # If distribution could not be found, call working_set.require
+    # to update the working set, and try to find the distribution
+    # again This might happen for e.g. when you install a package
+    # twice, once in editable mode and again in normal mode. Now
+    # when you want to run uninstall twice the package gets removed
+    # from the working set in the first uninstall which removes the
+    # normal package, so we have to populate the working set so that
+    # the packages gets picked up and is successfully uninstalled
+    # the second time too.
+    if not dist:
+        try:
+            pkg_resources.working_set.require(req_name)
+        except pkg_resources.DistributionNotFound:
+            return None
+    return search_distribution(req_name)
+
+
 def egg_link_path(dist):
     # type: (Distribution) -> Optional[str]
     """

tests/functional/test_search.py

@@ -168,7 +168,8 @@ def test_latest_prerelease_install_message(caplog, monkeypatch):
 
     dist = pretend.stub(version="1.0.0")
     get_dist = pretend.call_recorder(lambda x: dist)
-    monkeypatch.setattr("pip._vendor.pkg_resources.get_distribution", get_dist)
+    monkeypatch.setattr("pip._internal.commands.search.get_distribution",
+                        get_dist)
     with caplog.at_level(logging.INFO):
         print_results(hits)
 

tests/unit/test_self_check_outdated.py

@@ -7,7 +7,6 @@
 import pretend
 import pytest
 from mock import patch
-from pip._vendor import pkg_resources
 
 from pip._internal import self_outdated_check
 from pip._internal.models.candidate import InstallationCandidate
@@ -172,7 +171,7 @@ def test_pip_self_version_check(monkeypatch, stored_time, installed_ver,
                         pretend.call_recorder(lambda *a, **kw: None))
     monkeypatch.setattr(logger, 'debug',
                         pretend.call_recorder(lambda s, exc_info=None: None))
-    monkeypatch.setattr(pkg_resources, 'get_distribution',
+    monkeypatch.setattr(self_outdated_check, 'get_distribution',
                         lambda name: MockDistribution(installer))
 
     fake_state = pretend.stub(