allow targeted stealing (#9)

* allow targeted stealing

* bump to 0.4.0

Author: provinzkraut

README.md

@@ -46,12 +46,12 @@ jobs:
 
 Pytest-cdist comes with several CLI and pytest-ini options:
 
-| CLI                     | Ini                   | Allowed values                | Default |
-|-------------------------|-----------------------|-------------------------------|---------|
-| `--cdist-justify-items` | `cdist-justify-items` | `none`, `file`, `scope`       | `none`  |
-| `--cdist-group-steal`   | `--cdist-group-steal` | `<group number>:<percentage>` | -       |
-| `--cdist-report`        | -                     | -                             | false   |
-| `--cdist-report-dir`    | `cdist-report-dir`    |                               | `.`     |
+| CLI                     | Ini                   | Allowed values                                                               | Default |
+|-------------------------|-----------------------|------------------------------------------------------------------------------|---------|
+| `--cdist-justify-items` | `cdist-justify-items` | `none`, `file`, `scope`                                                      | `none`  |
+| `--cdist-group-steal`   | `--cdist-group-steal` | `<target group>:<percentage>` / `<target group>:<percentage>:<source group>` | -       |
+| `--cdist-report`        | -                     | -                                                                            | false   |
+| `--cdist-report-dir`    | `cdist-report-dir`    |                                                                              | `.`     |
 
 
 ### Controlling how items are split up
@@ -96,6 +96,15 @@ cdist-group-steal=2:30
 pytest --cdist-group=1/2 --cdist-group-steal=2:30
 ```
 
+It is also possible to redistribute items between two specific groups, by specifying 
+both as source and a target group. The following configuration would assign 50% of the 
+items in group 1 to group 2:
+
+```bash
+pytest --cdist-group=1/3 --cdist-group-steal=1:50:2
+```
+
+
 ### With pytest-xdist
 
 When running under pytest-xdist, pytest-cdist will honour tests marked with 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pytest-cdist"
-version = "0.3.1"
+version = "0.4.0"
 description = "A pytest plugin to split your test suite into multiple parts"
 authors = [
   { name = "Janek Nouvertné", email = "provinzkraut@posteo.de" },

pytest_cdist/plugin.py

@@ -5,6 +5,7 @@
 import json
 import os
 import pathlib
+import re
 from typing import TypeVar, Literal, TYPE_CHECKING
 
 import pytest
@@ -42,24 +43,76 @@ def _get_item_file(item: pytest.Item) -> str:
     return item.nodeid.split("::", 1)[0]
 
 
+@dataclasses.dataclass(frozen=True)
+class GroupStealOpt:
+    source_group: int | None
+    target_group: int
+    amount: int
+
+    def __str__(self) -> str:
+        out = f"g{self.target_group}:{self.amount}"
+        if self.source_group:
+            out += f":g{self.source_group}"
+        return out
+
+
+def _distribute(
+    groups: list[list[pytest.Item]],
+    from_: int,
+    to: int,
+    amount: int,
+) -> None:
+    items = groups[from_]
+    num_items_to_move = max(0, min(len(items), (len(items) * amount) // 100))
+    items_to_move = items[:num_items_to_move]
+    groups[to].extend(items_to_move)
+    groups[from_] = items[num_items_to_move:]
+
+
 def _distribute_with_bias(
-    groups: list[list[pytest.Item]], target: int, bias: int
+    groups: list[list[pytest.Item]],
+    group_steal: GroupStealOpt,
 ) -> list[list[pytest.Item]]:
-    for i, lst in enumerate(groups):
-        if i != target:
-            num_items_to_move = max(0, min(len(lst), (len(lst) * bias) // 100))
-            items_to_move = lst[:num_items_to_move]
-            groups[target].extend(items_to_move)
-            groups[i] = lst[num_items_to_move:]
+    source_groups = (
+        [group_steal.source_group]
+        if group_steal.source_group is not None
+        else range(len(groups))
+    )
+    for source_group in source_groups:
+        if source_group != group_steal.target_group:
+            _distribute(
+                groups,
+                from_=source_group,
+                to=group_steal.target_group,
+                amount=group_steal.amount,
+            )
 
     return groups
 
 
-def _get_group_steal_opt(opt: str | None) -> tuple[int, int] | None:
+def _get_group_steal_opt(opt: str | None) -> list[GroupStealOpt]:
     if opt is None:
-        return None
-    target_group, amount_to_steal = opt.split(":")
-    return int(target_group) - 1, int(amount_to_steal)
+        return []
+
+    opts = []
+    for group_opt in opt.split(","):
+        match = re.match(r"g?(\d+):(\d+)(?::g?(\d+))?", group_opt.strip())
+        if match is None:
+            raise ValueError(f"Invalid group steal option: {group_opt!r}")
+        source_group: int | None = None
+        target_group = int(match.group(1)) - 1
+        amount_to_steal = int(match.group(2))
+        if source_group_match := match.group(3):
+            source_group = int(source_group_match) - 1
+
+        opts.append(
+            GroupStealOpt(
+                source_group=source_group,
+                target_group=target_group,
+                amount=amount_to_steal,
+            )
+        )
+    return opts
 
 
 def _justify_items(
@@ -113,12 +166,12 @@ def _justify_xdist_groups(groups: list[list[pytest.Item]]) -> list[list[pytest.I
     return groups
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(kw_only=True)
 class CdistConfig:
     current_group: int
     total_groups: int
     justify_items_strategy: JustifyItemsStrategy = "none"
-    group_steal: tuple[int, int] | None = None
+    group_steal: list[GroupStealOpt]
     write_report: bool = False
     report_dir: pathlib.Path = pathlib.Path(".")
 
@@ -132,9 +185,8 @@ def cli_options(self, config: pytest.Config) -> str:
             opts.append(f"--cdist-justify-items={self.justify_items_strategy}")
 
         if self.group_steal and "cdist-group-steal" not in config.inicfg:
-            opts.append(
-                f"--cdist-group-steal={self.group_steal[0] + 1}:{self.group_steal[1]}"
-            )
+            steal_opt = ",".join(map(str, self.group_steal))
+            opts.append(f"--cdist-group-steal={steal_opt}")
 
         if self.write_report:
             opts.append("--cdist-report")
@@ -202,7 +254,7 @@ def pytest_addoption(parser: pytest.Parser) -> None:
         action="store",
         default=None,
         help="make a group steal a percentage of items from other groups. '1:30' would "
-        "make group 1 steal 30%% of items from all other groups)",
+        "make group 1 steal 30% of items from all other groups",
     )
 
     parser.addini("cdist-justify-items", help="justify items strategy", default="none")
@@ -238,12 +290,10 @@ def pytest_collection_modifyitems(
 
     groups = _partition_list(items, cdist_config.total_groups)
 
-    if cdist_config.group_steal is not None:
-        target_group, amount_to_steal = cdist_config.group_steal
+    for group_steal in cdist_config.group_steal:
         groups = _distribute_with_bias(
             groups,
-            target=target_group,
-            bias=amount_to_steal,
+            group_steal=group_steal,
         )
 
     if cdist_config.justify_items_strategy != "none":

tests/test_plugin.py

@@ -209,6 +209,84 @@ def test_four():
     result.assert_outcomes(passed=3, deselected=1)
 
 
+def test_steal_with_target(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile("""
+    def test_one():
+        assert False
+
+    def test_two():
+        assert True
+
+    def test_three():
+        assert True
+
+    def test_four():
+        assert True
+    """)
+
+    # natural distribution would be
+    # 1: test_one, test_two
+    # 2: test_three
+    # 3: test_four
+    # telling group 3 to steal 50% of group 1 should result in
+    # 1: test_two
+    # 2: test_three
+    # 3: test_four, test_one
+    cli_opt = "--cdist-group-steal=g3:50:g1"
+    result = pytester.runpytest_inprocess("--cdist-group=1/3", cli_opt)
+    result.assert_outcomes(passed=1, deselected=3)
+
+    result = pytester.runpytest_inprocess("--cdist-group=2/3", cli_opt)
+    result.assert_outcomes(passed=1, deselected=3)
+
+    result = pytester.runpytest_inprocess("--cdist-group=3/3", cli_opt)
+    result.assert_outcomes(passed=1, failed=1, deselected=2)
+
+
+def test_steal_multiple_target(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile("""
+    def test_one():
+        assert True
+
+    def test_two():
+        assert True
+
+    def test_three():
+        assert True
+
+    def test_four():
+        assert True
+        
+    def test_five():
+        assert True
+        
+    def test_six():
+        assert True
+    """)
+
+    # natural distribution would be
+    # 1: 2
+    # 2: 2
+    # 3: 2
+    # first, we're telling group 2 to steal 50% of all other groups:
+    # 1: 1
+    # 2: 4
+    # 3: 1
+    # then, we're telling group 3 to steal 50% of group 2
+    # 1: 1
+    # 2: 2
+    # 3: 3
+    cli_opt = "--cdist-group-steal=g2:50,g3:50:g2"
+    result = pytester.runpytest("--cdist-group=1/3", cli_opt)
+    result.assert_outcomes(passed=1, deselected=5)
+
+    result = pytester.runpytest("--cdist-group=2/3", cli_opt)
+    result.assert_outcomes(passed=2, deselected=4)
+
+    result = pytester.runpytest("--cdist-group=3/3", cli_opt)
+    result.assert_outcomes(passed=3, deselected=3)
+
+
 def test_steal_with_justify(pytester: pytest.Pytester) -> None:
     pytester.makepyfile("""
     class TestFoo: