From 00f953b38d8fd97fe1c17bfde1e378c1ff149d79 Mon Sep 17 00:00:00 2001
From: DetachHead <57028336+DetachHead@users.noreply.github.com>
Date: Sun, 2 Jun 2024 15:31:43 +1000
Subject: [PATCH] add hook to the pytest vscode plugin to determine how many
 xdist workers to use based on how many tests are selected

---
 python_files/vscode_pytest/__init__.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/python_files/vscode_pytest/__init__.py b/python_files/vscode_pytest/__init__.py
index 3534bf7c699f6..d4167115390f7 100644
--- a/python_files/vscode_pytest/__init__.py
+++ b/python_files/vscode_pytest/__init__.py
@@ -8,7 +8,7 @@
 import sys
 import traceback
 
-
+from pluggy import Result
 import pytest
 
 script_dir = pathlib.Path(__file__).parent.parent
@@ -16,7 +16,7 @@
 sys.path.append(os.fspath(script_dir / "lib" / "python"))
 
 from testing_tools import socket_manager  # noqa: E402
-from typing import Any, Dict, List, Optional, Union, TypedDict, Literal  # noqa: E402
+from typing import Any, Dict, Generator, List, Optional, Union, TypedDict, Literal  # noqa: E402
 
 
 class TestData(TypedDict):
@@ -881,3 +881,9 @@ def send_post_request(
             f"Plugin error, exception thrown while attempting to send data[vscode-pytest]: {error} \n[vscode-pytest] data: \n{data}\n",
             file=sys.stderr,
         )
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_xdist_auto_num_workers(config: pytest.Config) -> Generator[None, Result[int], int]:
+    """determine how many workers to use based on how many tests were selected in the test explorer"""
+    return min((yield).get_result(), len(config.option.file_or_dir))