From 219d82ff866ddf86da8fe78d02ba5fcba6e3b732 Mon Sep 17 00:00:00 2001
From: Philip Meier <github.pmeier@posteo.de>
Date: Thu, 13 Aug 2020 13:58:33 +0200
Subject: [PATCH] fix discovery in extra requirements (#8)

* fix discovery in extra requirements

* simplify
---
 tests/test_plugin.py | 37 +++++++++++++++++++++++++++++++++++--
 tox_ltt/plugin.py    | 13 +++++++++++--
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/tests/test_plugin.py b/tests/test_plugin.py
index 3b5cfad..086e92c 100644
--- a/tests/test_plugin.py
+++ b/tests/test_plugin.py
@@ -50,13 +50,17 @@ def get_setup_py():
     """
 
 
-def get_setup_cfg(name, version, install_requires=None):
+def get_setup_cfg(name, version, install_requires=None, extra_requires=None):
     lines = ["[metadata]", f"name = {name}", f"version = {version}"]
 
     if install_requires is not None:
         lines.extend(("[options]", "install_requires = "))
         lines.extend([f"\t{req}" for req in install_requires])
 
+    if extra_requires is not None:
+        lines.extend(("[options.extras_require]", "extra = "))
+        lines.extend([f"\t{req}" for req in extra_requires])
+
     return "\n".join(lines)
 
 
@@ -65,6 +69,7 @@ def get_tox_ini(
     force_cpu=None,
     deps=None,
     skip_install=False,
+    extra=False,
     pep517=True,
 ):
 
@@ -77,6 +82,8 @@ def get_tox_ini(
 
     if skip_install:
         lines.append("skip_install = True")
+    if extra:
+        lines.append("extras = extra")
     if disable_light_the_torch is not None:
         lines.append(f"disable_light_the_torch = {disable_light_the_torch}")
     if force_cpu is not None:
@@ -94,6 +101,7 @@ def tox_ltt_initproj_(
         name="foo",
         version="1.2.3",
         install_requires=None,
+        extra_requires=None,
         disable_light_the_torch=None,
         force_cpu=None,
         deps=None,
@@ -102,10 +110,14 @@ def tox_ltt_initproj_(
     ):
         filedefs = {
             "setup.cfg": get_setup_cfg(
-                name, version, install_requires=install_requires
+                name,
+                version,
+                install_requires=install_requires,
+                extra_requires=extra_requires,
             ),
             "tox.ini": get_tox_ini(
                 skip_install=skip_install,
+                extra=extra_requires is not None,
                 disable_light_the_torch=disable_light_the_torch,
                 force_cpu=force_cpu,
                 deps=deps,
@@ -236,3 +248,24 @@ def test_tox_ltt_project_pytorch_dists(
 
             args, _ = mock.call_args
             assert set(args[0]) == dists
+
+
+def test_tox_ltt_project_extra_pytorch_dists(
+    subtests, patch_find_links, tox_ltt_initproj, cmd, install_mock
+):
+    mock = patch_find_links()
+
+    extra_requires = ("torch>=1.5.0", "torchvision>=0.6.0")
+    dists = set(extra_requires)
+
+    for pep517 in (True, False):
+        mock.reset()
+        with subtests.test(pep517=pep517):
+            tox_ltt_initproj(extra_requires=extra_requires, pep517=pep517)
+
+            result = cmd()
+
+            result.assert_success(is_run_test_env=False)
+
+            args, _ = mock.call_args
+            assert set(args[0]) == dists
diff --git a/tox_ltt/plugin.py b/tox_ltt/plugin.py
index 1ca77b3..bfee7a0 100644
--- a/tox_ltt/plugin.py
+++ b/tox_ltt/plugin.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional, Sequence, cast
+from typing import Any, List, Optional, Sequence, cast
 
 import tox
 from tox import reporter
@@ -61,7 +61,10 @@ def tox_testenv_install_deps(venv: VirtualEnv, action: Action) -> None:
     requirements = [dep_config.name for dep_config in venv.get_resolved_dependencies()]
 
     if not envconfig.skip_install:
-        requirements.append(venv.package.strpath)
+        path = venv.package.strpath
+        if envconfig.extras:
+            path += f"[{','.join(envconfig.extras)}]"
+        requirements.append(path)
 
     if not requirements:
         return None
@@ -69,6 +72,7 @@ def tox_testenv_install_deps(venv: VirtualEnv, action: Action) -> None:
     action.setactivity("finddeps-light-the-torch", "")
 
     dists = ltt.extract_dists(requirements)
+    dists = remove_extras(dists)
 
     if not dists:
         reporter.verbosity1(
@@ -95,3 +99,8 @@ def tox_testenv_install_deps(venv: VirtualEnv, action: Action) -> None:
 
     action.setactivity("installdeps-light-the-torch", ", ".join(links))
     venv.run_install_command(links, action)
+
+
+# TODO: this should probably implemented in light-the-torch
+def remove_extras(dists: List[str]) -> List[str]:
+    return [dist.split(";")[0] for dist in dists]