Basic abstraction one actual migration

src/pip/_internal/metadata/__init__.py

@@ -0,0 +1,27 @@
+from pip._internal.utils.compat import lru_cache
+from pip._internal.utils.typing import MYPY_CHECK_RUNNING
+
+from .base import BaseDistribution, BaseEnvironment
+
+if MYPY_CHECK_RUNNING:
+    from typing import Optional, Sequence, Tuple
+
+
+__all__ = ["BaseDistribution", "BaseEnvironment", "get_environment"]
+
+
+@lru_cache(maxsize=None)
+def _get_environment(paths):
+    # type: (Optional[Tuple[str, ...]]) -> BaseEnvironment
+    from .pkg_resources import Environment
+    if paths is None:
+        return Environment.default()
+    return Environment.from_paths(paths)
+
+
+def get_environment(paths=None):
+    # type: (Optional[Sequence[str]]) -> BaseEnvironment
+    # Convert input to tuple since lru_cache() requires hashable arguments.
+    if paths is not None:
+        paths = tuple(paths)
+    return _get_environment(paths)

src/pip/_internal/metadata/base.py

@@ -0,0 +1,80 @@
+import abc
+
+from pip._vendor.six import add_metaclass
+
+from pip._internal.utils.misc import stdlib_pkgs  # TODO: Move definition here.
+from pip._internal.utils.typing import MYPY_CHECK_RUNNING
+
+if MYPY_CHECK_RUNNING:
+    from typing import Container, Iterator, Optional, Sequence
+
+
+@add_metaclass(abc.ABCMeta)
+class BaseDistribution(object):
+    @property
+    def canonical_name(self):
+        # type: () -> str
+        raise NotImplementedError()
+
+    @property
+    def installer(self):
+        # type: () -> str
+        raise NotImplementedError()
+
+    @property
+    def editable(self):
+        # type: () -> bool
+        raise NotImplementedError()
+
+    @property
+    def local(self):
+        # type: () -> bool
+        raise NotImplementedError()
+
+    @property
+    def in_usersite(self):
+        # type: () -> bool
+        raise NotImplementedError()
+
+
+@add_metaclass(abc.ABCMeta)
+class BaseEnvironment(object):
+    """An environment containing distributions to introspect.
+    """
+    @classmethod
+    def default(cls):
+        # type: () -> BaseEnvironment
+        raise NotImplementedError()
+
+    @classmethod
+    def from_paths(cls, paths):
+        # type: (Sequence[str]) -> BaseEnvironment
+        raise NotImplementedError()
+
+    def iter_distributions(self):
+        # type: () -> Iterator[BaseDistribution]
+        raise NotImplementedError()
+
+    def iter_installed_distributions(
+        self,
+        local_only=True,  # type: bool
+        skip=stdlib_pkgs,  # type: Container[str]
+        include_editables=True,  # type: bool
+        editables_only=False,  # type: bool
+        user_only=False,  # type: bool
+    ):
+        # type: (...) -> Iterator[BaseDistribution]
+        it = self.iter_distributions()
+        if local_only:
+            it = (d for d in it if d.local)
+        if not include_editables:
+            it = (d for d in it if not d.editable)
+        if editables_only:
+            it = (d for d in it if d.editable)
+        if user_only:
+            it = (d for d in it if d.in_usersite)
+        return (d for d in it if d.canonical_name not in skip)
+
+    def get_installed_distribution(self, name):
+        # type: (str) -> Optional[BaseDistribution]
+        raise NotImplementedError()

src/pip/_internal/metadata/pkg_resources.py

@@ -0,0 +1,94 @@
+from pip._vendor import pkg_resources
+from pip._vendor.packaging.utils import canonicalize_name
+
+from pip._internal.utils.misc import (
+    dist_in_usersite,
+    dist_is_editable,
+    dist_is_local,
+)
+from pip._internal.utils.packaging import get_installer
+from pip._internal.utils.typing import MYPY_CHECK_RUNNING
+
+from .base import BaseDistribution, BaseEnvironment
+
+if MYPY_CHECK_RUNNING:
+    from typing import Iterator, Optional, Sequence
+
+
+class Distribution(BaseDistribution):
+    def __init__(self, dist):
+        # type: (pkg_resources.Distribution) -> None
+        self._dist = dist
+
+    @property
+    def canonical_name(self):
+        # type: () -> str
+        return canonicalize_name(self._dist.project_name)
+
+    @property
+    def installer(self):
+        # type: () -> str
+        return get_installer(self._dist)
+
+    @property
+    def editable(self):
+        # type: () -> bool
+        return dist_is_editable(self._dist)
+
+    @property
+    def local(self):
+        # type: () -> bool
+        return dist_is_local(self._dist)
+
+    @property
+    def in_usersite(self):
+        # type: () -> bool
+        return dist_in_usersite(self._dist)
+
+
+class Environment(BaseEnvironment):
+    def __init__(self, ws):
+        # type: (pkg_resources.WorkingSet) -> None
+        self._ws = ws
+
+    @classmethod
+    def default(cls):
+        # type: () -> BaseEnvironment
+        return cls(pkg_resources.working_set)
+
+    @classmethod
+    def from_paths(cls, paths):
+        # type: (Sequence[str]) -> BaseEnvironment
+        return cls(pkg_resources.WorkingSet(paths))
+
+    def iter_distributions(self):
+        # type: () -> Iterator[BaseDistribution]
+        for dist in self._ws:
+            yield Distribution(dist)
+
+    def _get_installed_distribution_from_cache(self, name):
+        # type: (str) -> Optional[BaseDistribution]
+        name = canonicalize_name(name)
+        for dist in self.iter_installed_distributions():
+            if dist.canonical_name == name:
+                return dist
+        return None
+
+    def get_installed_distribution(self, name):
+        # type: (str) -> Optional[BaseDistribution]
+        dist = self._get_installed_distribution_from_cache(name)
+        if dist:
+            return dist
+        # If distribution could not be found, call WorkingSet.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
+        # using setup.py develop and again using setup.py install. Now when
+        # pip uninstall is run twice, the package gets removed from the
+        # working set in the first uninstall, so we have to populate the
+        # working set again so that pip knows about it and the packages gets
+        # picked up and is successfully uninstalled the second time too.
+        try:
+            self._ws.require(name)
+        except pkg_resources.DistributionNotFound:
+            return None
+        return self._get_installed_distribution_from_cache(name)

src/pip/_internal/utils/misc.py

@@ -18,7 +18,6 @@
 from collections import deque
 
 from pip._vendor import pkg_resources
-from pip._vendor.packaging.utils import canonicalize_name
 # 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.retrying import retry  # type: ignore
@@ -34,6 +33,7 @@
     site_packages,
     user_site,
 )
+from pip._internal.metadata import get_environment
 from pip._internal.utils.compat import (
     WINDOWS,
     expanduser,
@@ -58,6 +58,9 @@
     )
     from pip._vendor.pkg_resources import Distribution
 
+    from pip._internal.metadata import pkg_resources as metadata_pkg_resources
+
+    PipMetadataDistribution = metadata_pkg_resources.Distribution
     VersionInfo = Tuple[int, int, int]
 
 
@@ -413,6 +416,8 @@ def dist_is_editable(dist):
     return False
 
 
+# TODO: Refactor the call site to make it call get_environment() directly,
+# and operate on the wrapper PipMetadataDistribution instead.
 def get_installed_distributions(
         local_only=True,  # type: bool
         skip=stdlib_pkgs,  # type: Container[str]
@@ -441,78 +446,24 @@ def get_installed_distributions(
     If ``paths`` is set, only report the distributions present at the
     specified list of locations.
     """
-    if paths:
-        working_set = pkg_resources.WorkingSet(paths)
-    else:
-        working_set = pkg_resources.working_set
-
-    if local_only:
-        local_test = dist_is_local
-    else:
-        def local_test(d):
-            return True
-
-    if include_editables:
-        def editable_test(d):
-            return True
-    else:
-        def editable_test(d):
-            return not dist_is_editable(d)
-
-    if editables_only:
-        def editables_only_test(d):
-            return dist_is_editable(d)
-    else:
-        def editables_only_test(d):
-            return True
-
-    if user_only:
-        user_test = dist_in_usersite
-    else:
-        def user_test(d):
-            return True
-
-    return [d for d in working_set
-            if local_test(d) and
-            d.key not in skip and
-            editable_test(d) and
-            editables_only_test(d) and
-            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)
+    iterator = get_environment(paths).iter_installed_distributions(
+        local_only=local_only,
+        skip=skip,
+        include_editables=include_editables,
+        editables_only=editables_only,
+        user_only=user_only,
+    )
+    return [cast("PipMetadataDistribution", dist)._dist for dist in iterator]
 
 
+# TODO: Refactor the call site to make it call get_environment() directly,
+# and operate on the wrapper PipMetadataDistribution instead.
 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 using setup.py develop and again using setup.py install.
-    # Now when run pip uninstall twice, the package gets removed
-    # from the working set in the first uninstall, so we have to populate
-    # the working set again so that pip knows about it and the packages
-    # gets picked up and is successfully uninstalled the second time too.
+    dist = get_environment().get_installed_distribution(req_name)
     if not dist:
-        try:
-            pkg_resources.working_set.require(req_name)
-        except pkg_resources.DistributionNotFound:
-            return None
-    return search_distribution(req_name)
+        return None
+    casted = cast("PipMetadataDistribution", dist)
+    return casted._dist
 
 
 def egg_link_path(dist):