Add --source option to "poetry add" (#1912)

* Add --source option to 'poetry add'

* Add tests for 'poetry add --source'

poetry/console/commands/add.py

@@ -33,6 +33,12 @@ class AddCommand(EnvCommand, InitCommand):
             "Platforms for which the dependency must be installed.",
             flag=False,
         ),
+        option(
+            "source",
+            None,
+            "Name of the source to use to install the package.",
+            flag=False,
+        ),
         option("allow-prereleases", None, "Accept prereleases."),
         option(
             "dry-run",
@@ -86,7 +92,9 @@ If you do not specify a version constraint, poetry will choose a suitable one ba
                     raise ValueError("Package {} is already present".format(name))
 
         requirements = self._determine_requirements(
-            packages, allow_prereleases=self.option("allow-prereleases")
+            packages,
+            allow_prereleases=self.option("allow-prereleases"),
+            source=self.option("source"),
         )
 
         for _constraint in requirements:
@@ -123,6 +131,9 @@ If you do not specify a version constraint, poetry will choose a suitable one ba
             if self.option("platform"):
                 constraint["platform"] = self.option("platform")
 
+            if self.option("source"):
+                constraint["source"] = self.option("source")
+
             if len(constraint) == 1 and "version" in constraint:
                 constraint = constraint["version"]
 

poetry/console/commands/init.py

@@ -204,7 +204,7 @@ The <c1>init</c1> command creates a basic <comment>pyproject.toml</> file in the
             f.write(content)
 
     def _determine_requirements(
-        self, requires, allow_prereleases=False
+        self, requires, allow_prereleases=False, source=None
     ):  # type: (List[str], bool) -> List[Dict[str, str]]
         if not requires:
             requires = []
@@ -300,7 +300,9 @@ The <c1>init</c1> command creates a basic <comment>pyproject.toml</> file in the
             elif "version" not in requirement:
                 # determine the best version automatically
                 name, version = self._find_best_version_for_package(
-                    requirement["name"], allow_prereleases=allow_prereleases
+                    requirement["name"],
+                    allow_prereleases=allow_prereleases,
+                    source=source,
                 )
                 requirement["version"] = version
                 requirement["name"] = name
@@ -315,6 +317,7 @@ The <c1>init</c1> command creates a basic <comment>pyproject.toml</> file in the
                     requirement["name"],
                     requirement["version"],
                     allow_prereleases=allow_prereleases,
+                    source=source,
                 )
 
                 requirement["name"] = name
@@ -324,13 +327,13 @@ The <c1>init</c1> command creates a basic <comment>pyproject.toml</> file in the
         return result
 
     def _find_best_version_for_package(
-        self, name, required_version=None, allow_prereleases=False
+        self, name, required_version=None, allow_prereleases=False, source=None
     ):  # type: (...) -> Tuple[str, str]
         from poetry.version.version_selector import VersionSelector
 
         selector = VersionSelector(self._get_pool())
         package = selector.find_best_candidate(
-            name, required_version, allow_prereleases=allow_prereleases
+            name, required_version, allow_prereleases=allow_prereleases, source=source
         )
 
         if not package:

poetry/version/version_selector.py

@@ -15,6 +15,7 @@ class VersionSelector(object):
         package_name,  # type: str
         target_package_version=None,  # type:  Union[str, None]
         allow_prereleases=False,  # type: bool
+        source=None,  # type: str
     ):  # type: (...) -> Union[Package, bool]
         """
         Given a package name and optional version,
@@ -26,7 +27,7 @@ class VersionSelector(object):
             constraint = parse_constraint("*")
 
         candidates = self._pool.find_packages(
-            package_name, constraint, allow_prereleases=True
+            package_name, constraint, allow_prereleases=True, repository=source
         )
         only_prereleases = all([c.version.is_prerelease() for c in candidates])
 

tests/console/commands/test_add.py

@@ -4,6 +4,8 @@ import pytest
 
 from cleo.testers import CommandTester
 
+from poetry.repositories.legacy_repository import LegacyRepository
+from poetry.semver import Version
 from poetry.utils._compat import Path
 from tests.helpers import get_dependency
 from tests.helpers import get_package
@@ -634,6 +636,72 @@ Package operations: 1 install, 0 updates, 0 removals
     }
 
 
+def test_add_constraint_with_source(app, poetry, installer):
+    repo = LegacyRepository(name="my-index", url="https://my-index.fake")
+    repo.add_package(get_package("cachy", "0.2.0"))
+    repo._cache.store("matches").put("cachy:0.2.0", [Version.parse("0.2.0")], 5)
+
+    poetry.pool.add_repository(repo)
+
+    command = app.find("add")
+    tester = CommandTester(command)
+
+    tester.execute("cachy=0.2.0 --source my-index")
+
+    expected = """\
+
+Updating dependencies
+Resolving dependencies...
+
+Writing lock file
+
+
+Package operations: 1 install, 0 updates, 0 removals
+
+  - Installing cachy (0.2.0)
+"""
+
+    assert expected == tester.io.fetch_output()
+
+    assert len(installer.installs) == 1
+
+    content = app.poetry.file.read()["tool"]["poetry"]
+
+    assert "cachy" in content["dependencies"]
+    assert content["dependencies"]["cachy"] == {
+        "version": "0.2.0",
+        "source": "my-index",
+    }
+
+
+def test_add_constraint_with_source_that_does_not_exist(app):
+    command = app.find("add")
+    tester = CommandTester(command)
+
+    with pytest.raises(ValueError) as e:
+        tester.execute("foo --source i-dont-exist")
+
+    assert 'Repository "i-dont-exist" does not exist.' == str(e.value)
+
+
+def test_add_constraint_not_found_with_source(app, poetry, mocker):
+    repo = LegacyRepository(name="my-index", url="https://my-index.fake")
+    mocker.patch.object(repo, "find_packages", return_value=[])
+
+    poetry.pool.add_repository(repo)
+
+    pypi = poetry.pool.repositories[0]
+    pypi.add_package(get_package("cachy", "0.2.0"))
+
+    command = app.find("add")
+    tester = CommandTester(command)
+
+    with pytest.raises(ValueError) as e:
+        tester.execute("cachy --source my-index")
+
+    assert "Could not find a matching version of package cachy" == str(e.value)
+
+
 def test_add_to_section_that_does_no_exist_yet(app, repo, installer):
     command = app.find("add")
     tester = CommandTester(command)