Refactor collect_requirements function signature to take general return values

Author: connesy

tests/helpers.py

@@ -3,7 +3,7 @@
 import sys
 from functools import wraps
 from pathlib import Path
-from typing import Callable, Optional
+from typing import Callable, Optional, TypeVar
 
 from tests.types import P, RawFilesDict, RequirementsDict, RequirementsStem
 
@@ -53,14 +53,15 @@ def run_command(
     result.check_returncode()
 
 
-def collect_requirements(
-    func: Callable[P, tuple[RawFilesDict, RequirementsStem]]
-) -> Callable[P, tuple[RequirementsDict, RequirementsStem]]:
+R = TypeVar("R")
+
+
+def collect_requirements(func: Callable[P, tuple[RawFilesDict, R]]) -> Callable[P, tuple[RequirementsDict, R]]:
     @wraps(func)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> tuple[RequirementsDict, RequirementsStem]:
-        files_dict, requirements_stem = func(*args, **kwargs)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> tuple[RequirementsDict, R]:
+        files_dict, other = func(*args, **kwargs)
         requirements_dict = {filename: "\n".join(requirements) for filename, requirements in files_dict.items()}
-        return requirements_dict, requirements_stem
+        return requirements_dict, other
 
     return wrapper