From 3f1cdd3ea548890466e06249d82e029dc45d79d7 Mon Sep 17 00:00:00 2001
From: Tzu-ping Chung <uranusjr@gmail.com>
Date: Thu, 23 Apr 2020 12:01:05 +0800
Subject: [PATCH] Basic abstraction one actual migration

---
 src/pip/_internal/metadata/__init__.py      |  16 +++
 src/pip/_internal/metadata/base.py          |  80 +++++++++++
 src/pip/_internal/metadata/pkg_resources.py |  94 +++++++++++++
 src/pip/_internal/self_outdated_check.py    |  14 +-
 src/pip/_internal/utils/misc.py             |  85 +++---------
 tests/unit/test_self_check_outdated.py      |  12 +-
 tests/unit/test_utils.py                    | 141 +++++++++-----------
 7 files changed, 281 insertions(+), 161 deletions(-)
 create mode 100644 src/pip/_internal/metadata/__init__.py
 create mode 100644 src/pip/_internal/metadata/base.py
 create mode 100644 src/pip/_internal/metadata/pkg_resources.py

diff --git a/src/pip/_internal/metadata/__init__.py b/src/pip/_internal/metadata/__init__.py
new file mode 100644
index 00000000000..40da099bd2b
--- /dev/null
+++ b/src/pip/_internal/metadata/__init__.py
@@ -0,0 +1,16 @@
+from pip._internal.utils.compat import lru_cache
+from pip._internal.utils.typing import MYPY_CHECK_RUNNING
+
+if MYPY_CHECK_RUNNING:
+    from typing import List, Optional
+
+    from .base import BaseEnvironment
+
+
+@lru_cache(maxsize=None)
+def get_environment(paths=None):
+    # type: (Optional[List[str]]) -> BaseEnvironment
+    from .pkg_resources import Environment
+    if paths is None:
+        return Environment.default()
+    return Environment.from_paths(paths)
diff --git a/src/pip/_internal/metadata/base.py b/src/pip/_internal/metadata/base.py
new file mode 100644
index 00000000000..fabe5c09296
--- /dev/null
+++ b/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, List, Optional
+
+
+@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: (List[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()
diff --git a/src/pip/_internal/metadata/pkg_resources.py b/src/pip/_internal/metadata/pkg_resources.py
new file mode 100644
index 00000000000..d7b949f4f33
--- /dev/null
+++ b/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, List, Optional
+
+
+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: (List[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)
diff --git a/src/pip/_internal/self_outdated_check.py b/src/pip/_internal/self_outdated_check.py
index fbd9dfd48b7..aaf4f079eb3 100644
--- a/src/pip/_internal/self_outdated_check.py
+++ b/src/pip/_internal/self_outdated_check.py
@@ -12,18 +12,14 @@
 
 from pip._internal.index.collector import LinkCollector
 from pip._internal.index.package_finder import PackageFinder
+from pip._internal.metadata import get_environment
 from pip._internal.models.selection_prefs import SelectionPreferences
 from pip._internal.utils.filesystem import (
     adjacent_tmp_file,
     check_path_owner,
     replace,
 )
-from pip._internal.utils.misc import (
-    ensure_dir,
-    get_distribution,
-    get_installed_version,
-)
-from pip._internal.utils.packaging import get_installer
+from pip._internal.utils.misc import ensure_dir, get_installed_version
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 
 if MYPY_CHECK_RUNNING:
@@ -113,10 +109,8 @@ 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.
     """
-    dist = get_distribution(pkg)
-    if not dist:
-        return False
-    return "pip" == get_installer(dist)
+    dist = get_environment().get_installed_distribution(pkg)
+    return dist is not None and "pip" == dist.installer
 
 
 def pip_self_version_check(session, options):
diff --git a/src/pip/_internal/utils/misc.py b/src/pip/_internal/utils/misc.py
index 7cf9944fdd3..2700a2b061a 100644
--- a/src/pip/_internal/utils/misc.py
+++ b/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]
 
 
@@ -441,78 +444,22 @@ 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]
 
 
 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):
diff --git a/tests/unit/test_self_check_outdated.py b/tests/unit/test_self_check_outdated.py
index c5e60d92fc4..daad7db5549 100644
--- a/tests/unit/test_self_check_outdated.py
+++ b/tests/unit/test_self_check_outdated.py
@@ -57,6 +57,14 @@ def get_metadata_lines(self, name):
             raise NotImplementedError('nope')
 
 
+class MockEnvironment(object):
+    def __init__(self, installer):
+        self.installer = installer
+
+    def get_installed_distribution(self, name):
+        return MockDistribution(self.installer)
+
+
 def _options():
     ''' Some default options that we pass to
     self_outdated_check.pip_self_version_check '''
@@ -97,8 +105,8 @@ 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(self_outdated_check, 'get_distribution',
-                        lambda name: MockDistribution(installer))
+    monkeypatch.setattr(self_outdated_check, 'get_environment',
+                        lambda: MockEnvironment(installer))
 
     fake_state = pretend.stub(
         state={"last_check": stored_time, 'pypi_version': installed_ver},
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index ebabd29e260..ea4d9f355ae 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -20,6 +20,7 @@
     HashMissing,
     InstallationError,
 )
+from pip._internal.metadata.base import BaseEnvironment
 from pip._internal.utils.deprecation import PipDeprecationWarning, deprecated
 from pip._internal.utils.encoding import BOMS, auto_decode
 from pip._internal.utils.glibc import (
@@ -189,100 +190,80 @@ def test_noegglink_in_sitepkgs_venv_global(self):
         assert egg_link_path(self.mock_dist) is None
 
 
-@patch('pip._internal.utils.misc.dist_in_usersite')
-@patch('pip._internal.utils.misc.dist_is_local')
-@patch('pip._internal.utils.misc.dist_is_editable')
+class MockDistribution(object):
+    def __init__(
+        self,
+        canonical_name,
+        site="local",
+        local=True,
+        in_usersite=False,
+        editable=False,
+    ):
+        assert site in ("local", "global", "user")
+        self._dist = self
+        self.canonical_name = canonical_name
+        self.local = site == "local"
+        self.in_usersite = site == "user"
+        self.editable = editable
+
+
 class Tests_get_installed_distributions:
     """test util.get_installed_distributions"""
 
-    workingset = [
-        Mock(test_name="global"),
-        Mock(test_name="editable"),
-        Mock(test_name="normal"),
-        Mock(test_name="user"),
-    ]
-
-    workingset_stdlib = [
-        Mock(test_name='normal', key='argparse'),
-        Mock(test_name='normal', key='wsgiref')
-    ]
-
-    workingset_freeze = [
-        Mock(test_name='normal', key='pip'),
-        Mock(test_name='normal', key='setuptools'),
-        Mock(test_name='normal', key='distribute')
-    ]
-
-    def dist_is_editable(self, dist):
-        return dist.test_name == "editable"
-
-    def dist_is_local(self, dist):
-        return dist.test_name != "global" and dist.test_name != 'user'
-
-    def dist_in_usersite(self, dist):
-        return dist.test_name == "user"
-
-    @patch('pip._vendor.pkg_resources.working_set', workingset)
-    def test_editables_only(self, mock_dist_is_editable,
-                            mock_dist_is_local,
-                            mock_dist_in_usersite):
-        mock_dist_is_editable.side_effect = self.dist_is_editable
-        mock_dist_is_local.side_effect = self.dist_is_local
-        mock_dist_in_usersite.side_effect = self.dist_in_usersite
+    class MockEnvironment(BaseEnvironment):
+        def iter_distributions(self):
+            yield MockDistribution('global', site="global")
+            yield MockDistribution('editable', editable=True)
+            yield MockDistribution('normal')
+            yield MockDistribution('user', site="user")
+
+    class MockEnvironmentStdlib(BaseEnvironment):
+        def iter_distributions(self):
+            yield MockDistribution('argparse')
+            yield MockDistribution('wsgiref')
+
+    class MockEnvironmentFreeze(BaseEnvironment):
+        def iter_distributions(self):
+            yield MockDistribution('pip')
+            yield MockDistribution('setuptools')
+            yield MockDistribution('distribute')
+
+    @patch('pip._internal.utils.misc.get_environment',
+           return_value=MockEnvironment())
+    def test_editables_only(self, mock_get_environment):
         dists = get_installed_distributions(editables_only=True)
         assert len(dists) == 1, dists
-        assert dists[0].test_name == "editable"
-
-    @patch('pip._vendor.pkg_resources.working_set', workingset)
-    def test_exclude_editables(self, mock_dist_is_editable,
-                               mock_dist_is_local,
-                               mock_dist_in_usersite):
-        mock_dist_is_editable.side_effect = self.dist_is_editable
-        mock_dist_is_local.side_effect = self.dist_is_local
-        mock_dist_in_usersite.side_effect = self.dist_in_usersite
+        assert dists[0].canonical_name == "editable"
+
+    @patch('pip._internal.utils.misc.get_environment',
+           return_value=MockEnvironment())
+    def test_exclude_editables(self, mock_get_environment):
         dists = get_installed_distributions(include_editables=False)
         assert len(dists) == 1
-        assert dists[0].test_name == "normal"
-
-    @patch('pip._vendor.pkg_resources.working_set', workingset)
-    def test_include_globals(self, mock_dist_is_editable,
-                             mock_dist_is_local,
-                             mock_dist_in_usersite):
-        mock_dist_is_editable.side_effect = self.dist_is_editable
-        mock_dist_is_local.side_effect = self.dist_is_local
-        mock_dist_in_usersite.side_effect = self.dist_in_usersite
+        assert dists[0].canonical_name == "normal"
+
+    @patch('pip._internal.utils.misc.get_environment',
+           return_value=MockEnvironment())
+    def test_include_globals(self, mock_get_environment):
         dists = get_installed_distributions(local_only=False)
         assert len(dists) == 4
 
-    @patch('pip._vendor.pkg_resources.working_set', workingset)
-    def test_user_only(self, mock_dist_is_editable,
-                       mock_dist_is_local,
-                       mock_dist_in_usersite):
-        mock_dist_is_editable.side_effect = self.dist_is_editable
-        mock_dist_is_local.side_effect = self.dist_is_local
-        mock_dist_in_usersite.side_effect = self.dist_in_usersite
-        dists = get_installed_distributions(local_only=False,
-                                            user_only=True)
+    @patch('pip._internal.utils.misc.get_environment',
+           return_value=MockEnvironment())
+    def test_user_only(self, mock_get_environment):
+        dists = get_installed_distributions(local_only=False, user_only=True)
         assert len(dists) == 1
-        assert dists[0].test_name == "user"
-
-    @patch('pip._vendor.pkg_resources.working_set', workingset_stdlib)
-    def test_gte_py27_excludes(self, mock_dist_is_editable,
-                               mock_dist_is_local,
-                               mock_dist_in_usersite):
-        mock_dist_is_editable.side_effect = self.dist_is_editable
-        mock_dist_is_local.side_effect = self.dist_is_local
-        mock_dist_in_usersite.side_effect = self.dist_in_usersite
+        assert dists[0].canonical_name == "user"
+
+    @patch('pip._internal.utils.misc.get_environment',
+           return_value=MockEnvironmentStdlib())
+    def test_gte_py27_excludes(self, mock_get_environment):
         dists = get_installed_distributions()
         assert len(dists) == 0
 
-    @patch('pip._vendor.pkg_resources.working_set', workingset_freeze)
-    def test_freeze_excludes(self, mock_dist_is_editable,
-                             mock_dist_is_local,
-                             mock_dist_in_usersite):
-        mock_dist_is_editable.side_effect = self.dist_is_editable
-        mock_dist_is_local.side_effect = self.dist_is_local
-        mock_dist_in_usersite.side_effect = self.dist_in_usersite
+    @patch('pip._internal.utils.misc.get_environment',
+           return_value=MockEnvironmentFreeze())
+    def test_freeze_excludes(self, mock_get_environment):
         dists = get_installed_distributions(
             skip=('setuptools', 'pip', 'distribute'))
         assert len(dists) == 0