From c2adb3283c0b76e290a8f25053bf035500b47cb2 Mon Sep 17 00:00:00 2001
From: Arun Babu Neelicattu <arun.neelicattu@gmail.com>
Date: Wed, 7 Oct 2020 15:51:26 +0200
Subject: [PATCH] locker: move export logic into locker for reuse

This change, moves common functionality from exporter into the locker
to allow for sharing of logic that can generate `DependencyPackage`
instances given a a list of requirements using the lock data.
---
 poetry/packages/locker.py | 44 +++++++++++++++++++++++++++++++++++++--
 poetry/utils/exporter.py  | 34 +++++++-----------------------
 2 files changed, 50 insertions(+), 28 deletions(-)

diff --git a/poetry/packages/locker.py b/poetry/packages/locker.py
index 2283a2373d2..a486bf1c5d6 100644
--- a/poetry/packages/locker.py
+++ b/poetry/packages/locker.py
@@ -6,9 +6,11 @@
 
 from copy import deepcopy
 from hashlib import sha256
-from typing import Any
+from typing import Iterator
 from typing import List
 from typing import Optional
+from typing import Sequence
+from typing import Union
 
 from tomlkit import array
 from tomlkit import document
@@ -25,8 +27,10 @@
 from poetry.core.semver.version import Version
 from poetry.core.toml.file import TOMLFile
 from poetry.core.version.markers import parse_marker
+from poetry.packages import DependencyPackage
 from poetry.utils._compat import OrderedDict
 from poetry.utils._compat import Path
+from poetry.utils.extras import get_extra_package_names
 
 
 logger = logging.getLogger(__name__)
@@ -183,7 +187,7 @@ def locked_repository(
 
     def get_project_dependencies(
         self, project_requires, pinned_versions=False, with_nested=False, with_dev=False
-    ):  # type: (List[Dependency], bool, bool, bool) -> Any
+    ):  # type: (List[Dependency], bool, bool, bool) -> List[Dependency]
         packages = self.locked_repository(with_dev).packages
 
         # group packages entries by name, this is required because requirement might use
@@ -263,6 +267,42 @@ def __get_locked_package(
             key=lambda x: x.name.lower(),
         )
 
+    def get_project_dependency_packages(
+        self, project_requires, dev=False, extras=None
+    ):  # type: (List[Dependency], bool, Optional[Union[bool, Sequence[str]]]) -> Iterator[DependencyPackage]
+        repository = self.locked_repository(with_dev_reqs=dev)
+
+        # Build a set of all packages required by our selected extras
+        extra_package_names = (
+            None if (isinstance(extras, bool) and extras is True) else ()
+        )
+
+        if extra_package_names is not None:
+            extra_package_names = set(
+                get_extra_package_names(
+                    repository.packages, self.lock_data.get("extras", {}), extras or (),
+                )
+            )
+
+        for dependency in self.get_project_dependencies(
+            project_requires=project_requires, with_nested=True, with_dev=dev,
+        ):
+            try:
+                package = repository.find_packages(dependency=dependency)[0]
+            except IndexError:
+                continue
+
+            # If a package is optional and we haven't opted in to it, continue
+            if extra_package_names is not None and (
+                package.optional and package.name not in extra_package_names
+            ):
+                continue
+
+            for extra in dependency.extras:
+                package.requires_extras.append(extra)
+
+            yield DependencyPackage(dependency=dependency, package=package)
+
     def set_lock_data(self, root, packages):  # type: (...) -> bool
         files = table()
         packages = self._lock_packages(packages)
diff --git a/poetry/utils/exporter.py b/poetry/utils/exporter.py
index 4200abd146d..ae85c29cc0b 100644
--- a/poetry/utils/exporter.py
+++ b/poetry/utils/exporter.py
@@ -1,3 +1,5 @@
+from typing import Optional
+from typing import Sequence
 from typing import Union
 
 from clikit.api.io import IO
@@ -5,7 +7,6 @@
 from poetry.poetry import Poetry
 from poetry.utils._compat import Path
 from poetry.utils._compat import decode
-from poetry.utils.extras import get_extra_package_names
 
 
 class Exporter(object):
@@ -51,38 +52,19 @@ def _export_requirements_txt(
         dev=False,
         extras=None,
         with_credentials=False,
-    ):  # type: (Path, Union[IO, str], bool, bool, bool) -> None
+    ):  # type: (Path, Union[IO, str], bool, bool, Optional[Union[bool, Sequence[str]]], bool) -> None
         indexes = set()
         content = ""
-        repository = self._poetry.locker.locked_repository(dev)
-
-        # Build a set of all packages required by our selected extras
-        extra_package_names = set(
-            get_extra_package_names(
-                repository.packages,
-                self._poetry.locker.lock_data.get("extras", {}),
-                extras or (),
-            )
-        )
-
         dependency_lines = set()
 
-        for dependency in self._poetry.locker.get_project_dependencies(
-            project_requires=self._poetry.package.all_requires,
-            with_nested=True,
-            with_dev=dev,
+        for dependency_package in self._poetry.locker.get_project_dependency_packages(
+            project_requires=self._poetry.package.all_requires, dev=dev, extras=extras
         ):
-            try:
-                package = repository.find_packages(dependency=dependency)[0]
-            except IndexError:
-                continue
-
-            # If a package is optional and we haven't opted in to it, continue
-            if package.optional and package.name not in extra_package_names:
-                continue
-
             line = ""
 
+            dependency = dependency_package.dependency
+            package = dependency_package.package
+
             if package.develop:
                 line += "-e "