feat: improve xdist compatibility

Author: epenet

src/syrupy/data.py

@@ -4,6 +4,7 @@
 )
 from typing import (
     TYPE_CHECKING,
+    Any,
     Dict,
     Iterator,
     List,
@@ -124,6 +125,16 @@ def __iter__(self) -> Iterator["SnapshotCollection"]:
     def __contains__(self, key: str) -> bool:
         return key in self._snapshot_collections
 
+    def serialize(self) -> dict[str, Any]:
+        return {k: [c.name for c in v] for k, v in self._snapshot_collections.items()}
+
+    def merge_serialized(self, data: dict[str, Any]) -> None:
+        for location, names in data.items():
+            snapshot_collection = SnapshotCollection(location=location)
+            for name in names:
+                snapshot_collection.add(Snapshot(name))
+            self.update(snapshot_collection)
+
 
 @dataclass
 class DiffedLine:

src/syrupy/report.py

@@ -508,6 +508,37 @@ def _ran_items_match_location(self, snapshot_location: str) -> bool:
             for item in self.ran_items
         )
 
+    def serialize(self) -> dict[str, Any]:
+        return {
+            "discovered": self.discovered.serialize(),
+            "created": self.created.serialize(),
+            "failed": self.failed.serialize(),
+            "matched": self.matched.serialize(),
+            "updated": self.updated.serialize(),
+            "used": self.used.serialize(),
+            "_collected_items": [
+                {
+                    "nodeid": c.nodeid,
+                    "name": c.name,
+                    "path": str(c.path),
+                    "modulename": c.obj.__module__,  # type: ignore[attr-defined]
+                    "methodname": c.obj.__name__,  # type: ignore[attr-defined]
+                }
+                for c in list(self.collected_items)
+            ],
+            "_selected_items": {
+                key: status.value for key, status in self.selected_items.items()
+            },
+        }
+
+    def merge_serialized(self, data: dict[str, Any]) -> None:
+        self.discovered.merge_serialized(data["discovered"])
+        self.created.merge_serialized(data["created"])
+        self.failed.merge_serialized(data["failed"])
+        self.matched.merge_serialized(data["matched"])
+        self.updated.merge_serialized(data["updated"])
+        self.used.merge_serialized(data["used"])
+
 
 @dataclass(frozen=True)
 class Expression:

src/syrupy/session.py

@@ -1,3 +1,5 @@
+import json
+import os
 from collections import defaultdict
 from dataclasses import (
     dataclass,
@@ -46,6 +48,20 @@ class ItemStatus(Enum):
     SKIPPED = "skipped"
 
 
+class _FakePytestObject:
+    def __init__(self, collected_item: dict[str, str]) -> None:
+        self.__module__ = collected_item["modulename"]
+        self.__name__ = collected_item["methodname"]
+
+
+class _FakePytestItem:
+    def __init__(self, collected_item: dict[str, str]) -> None:
+        self.nodeid = collected_item["nodeid"]
+        self.name = collected_item["name"]
+        self.path = Path(collected_item["path"])
+        self.obj = _FakePytestObject(collected_item)
+
+
 @dataclass
 class SnapshotSession:
     pytest_session: "pytest.Session"
@@ -127,6 +143,24 @@ def ran_item(
             except ValueError:
                 pass  # if we don't understand the outcome, leave the item as "not run"
 
+    def _merge_collected_items(self, collected_items: list[dict[str, str]]) -> None:
+        for collected_item in collected_items:
+            custom_item = _FakePytestItem(collected_item)
+            if not any(
+                t.nodeid == custom_item.nodeid and t.name == custom_item.nodeid
+                for t in self._collected_items
+            ):
+                self._collected_items.add(custom_item)  # type: ignore[arg-type]
+
+    def _merge_selected_items(self, selected_items: dict[str, str]) -> None:
+        for key, selected_item in selected_items.items():
+            if key in self._selected_items:
+                status = ItemStatus(selected_item)
+                if status != ItemStatus.NOT_RUN:
+                    self._selected_items[key] = status
+            else:
+                self._selected_items[key] = ItemStatus(selected_item)
+
     def finish(self) -> int:
         exitstatus = 0
         self.flush_snapshot_write_queue()
@@ -139,16 +173,39 @@ def finish(self) -> int:
         )
 
         if is_xdist_worker():
-            # TODO: If we're in a pytest-xdist worker, we need to combine the reports
-            # of all the workers so that the controller can handle unused
-            # snapshot removal.
+            worker_count = os.getenv("PYTEST_XDIST_WORKER_COUNT")
+            with open(".pytest_syrupy_worker_count", "w", encoding="utf-8") as f:
+                f.write(worker_count)  # type: ignore[arg-type]
+            with open(
+                f".pytest_syrupy_{os.getenv("PYTEST_XDIST_WORKER")}_result",
+                "w",
+                encoding="utf-8",
+            ) as f:
+                json.dump(self.report.serialize(), f, indent=2)
             return exitstatus
         elif is_xdist_controller():
             # TODO: If we're in a pytest-xdist controller, merge all the reports.
             # Until this is implemented, running syrupy with pytest-xdist is only
             # partially functional.
             return exitstatus
 
+        worker_count = None
+        try:
+            with open(".pytest_syrupy_worker_count", encoding="utf-8") as f:
+                worker_count = f.read()
+            os.remove(".pytest_syrupy_worker_count")
+        except FileNotFoundError:
+            pass
+
+        if worker_count:
+            for i in range(int(worker_count)):
+                with open(f".pytest_syrupy_gw{i}_result", encoding="utf-8") as f:
+                    data = json.load(f)
+                    self._merge_collected_items(data["_collected_items"])
+                    self._merge_selected_items(data["_selected_items"])
+                    self.report.merge_serialized(data)
+                os.remove(f".pytest_syrupy_gw{i}_result")
+
         if self.report.num_unused:
             if self.update_snapshots:
                 self.remove_unused_snapshots(