From 42c19ec9bed96cebc4dab284d9da28be15c8ccd2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nikit=D0=B0=20Tsvetk=D0=BEv?= Date: Sun, 14 Jan 2024 20:18:12 +0400 Subject: [PATCH] add experimental selective discoverer (#78) --- tests/plugins/skipper/_utils.py | 3 +- vedro/plugins/skipper/_discoverer.py | 39 ++++++++++++++++++++++++ vedro/plugins/skipper/_skipper.py | 45 +++++++++++++++++++++++++--- 3 files changed, 82 insertions(+), 5 deletions(-) create mode 100644 vedro/plugins/skipper/_discoverer.py diff --git a/tests/plugins/skipper/_utils.py b/tests/plugins/skipper/_utils.py index dd1a9280..ce9fa3a1 100644 --- a/tests/plugins/skipper/_utils.py +++ b/tests/plugins/skipper/_utils.py @@ -45,7 +45,8 @@ async def fire_arg_parsed_event(dispatcher: Dispatcher, *, file_or_dir = ["."] if ignore is None: ignore = [] - namespace = Namespace(subject=subject, file_or_dir=file_or_dir, ignore=ignore) + namespace = Namespace(subject=subject, file_or_dir=file_or_dir, ignore=ignore, + exp_selective_discoverer=False) arg_parsed_event = ArgParsedEvent(namespace) await dispatcher.fire(arg_parsed_event) diff --git a/vedro/plugins/skipper/_discoverer.py b/vedro/plugins/skipper/_discoverer.py new file mode 100644 index 00000000..e8f7e0ca --- /dev/null +++ b/vedro/plugins/skipper/_discoverer.py @@ -0,0 +1,39 @@ +from pathlib import Path +from typing import List, Optional, Set + +from vedro.core import ScenarioFinder, ScenarioLoader, ScenarioOrderer, VirtualScenario +from vedro.core.scenario_discoverer import ScenarioDiscoverer, create_vscenario + +__all__ = ("SelectiveScenarioDiscoverer",) + + +class SelectiveScenarioDiscoverer(ScenarioDiscoverer): + def __init__(self, + finder: ScenarioFinder, loader: ScenarioLoader, orderer: ScenarioOrderer, *, + selected_paths: Optional[Set[Path]] = None) -> None: + super().__init__(finder, loader, orderer) + self._selected_paths = selected_paths + + async def discover(self, root: Path) -> List[VirtualScenario]: + scenarios = [] + async for path in self._finder.find(root): + if not self._is_path_selected(path): + continue + loaded = await self._loader.load(path) + for scn in loaded: + scenarios.append(create_vscenario(scn)) + return await self._orderer.sort(scenarios) + + def _is_path_selected(self, path: Path) -> bool: + if self._selected_paths is None or len(self._selected_paths) == 0: + return True + abs_path = path.absolute() + return any(self._is_relative_to(abs_path, selected) for selected in self._selected_paths) + + def _is_relative_to(self, path: Path, parent: Path) -> bool: + try: + path.relative_to(parent) + except ValueError: + return False + else: + return True diff --git a/vedro/plugins/skipper/_skipper.py b/vedro/plugins/skipper/_skipper.py index 2ee84268..c3cbebe9 100644 --- a/vedro/plugins/skipper/_skipper.py +++ b/vedro/plugins/skipper/_skipper.py @@ -1,8 +1,11 @@ import os -from typing import Any, List, Optional, Type, Union, cast +from pathlib import Path +from typing import Any, List, Optional, Set, Type, Union, cast -from vedro.core import Dispatcher, Plugin, PluginConfig, VirtualScenario -from vedro.events import ArgParsedEvent, ArgParseEvent, StartupEvent +from vedro.core import ConfigType, Dispatcher, Plugin, PluginConfig, VirtualScenario +from vedro.events import ArgParsedEvent, ArgParseEvent, ConfigLoadedEvent, StartupEvent + +from ._discoverer import SelectiveScenarioDiscoverer as ScenarioDiscoverer __all__ = ("Skipper", "SkipperPlugin",) @@ -17,16 +20,21 @@ def __init__(self, file_path: str, cls_name: Optional[str], tmpl_idx: Optional[i class SkipperPlugin(Plugin): def __init__(self, config: Type["Skipper"]) -> None: super().__init__(config) + self._global_config: Union[ConfigType, None] = None self._subject: Union[str, None] = None self._selected: List[_CompositePath] = [] self._deselected: List[_CompositePath] = [] self._forbid_only = config.forbid_only def subscribe(self, dispatcher: Dispatcher) -> None: - dispatcher.listen(ArgParseEvent, self.on_arg_parse) \ + dispatcher.listen(ConfigLoadedEvent, self.on_config_loaded) \ + .listen(ArgParseEvent, self.on_arg_parse) \ .listen(ArgParsedEvent, self.on_arg_parsed) \ .listen(StartupEvent, self.on_startup) + def on_config_loaded(self, event: ConfigLoadedEvent) -> None: + self._global_config = event.config + def on_arg_parse(self, event: ArgParseEvent) -> None: event.arg_parser.add_argument("file_or_dir", nargs="*", default=["."], help="Select scenarios in a given file or directory") @@ -34,6 +42,16 @@ def on_arg_parse(self, event: ArgParseEvent) -> None: help="Skip scenarios in a given file or directory") event.arg_parser.add_argument("--subject", help="Select scenarios with a given subject") + help_message = ( + "Enables the experimental selective discoverer feature, " + "which optimizes startup speed by loading scenarios only from specified files. " + "This is particularly beneficial for very large test suites " + "where Python's import mechanism can be slow, " + "thus reducing the initial load time and improving overall test execution efficiency." + ) + event.arg_parser.add_argument("--exp-selective-discoverer", action="store_true", + help=help_message) + def on_arg_parsed(self, event: ArgParsedEvent) -> None: self._subject = event.args.subject @@ -49,6 +67,25 @@ def on_arg_parsed(self, event: ArgParsedEvent) -> None: assert os.path.isdir(path) or os.path.isfile(path), f"{path!r} does not exist" self._deselected.append(composite_path) + exp_selective_discoverer = event.args.exp_selective_discoverer + if exp_selective_discoverer and len(self._deselected) == 0 and self._subject is None: + assert self._global_config is not None # for type checking + self._global_config.Registry.ScenarioDiscoverer.register(lambda: ScenarioDiscoverer( + finder=self._global_config.Registry.ScenarioFinder(), + loader=self._global_config.Registry.ScenarioLoader(), + orderer=self._global_config.Registry.ScenarioOrderer(), + selected_paths=self.__get_selected_paths(), + ), self) + + def __get_selected_paths(self) -> Set[Path]: + selected_paths = set() + default_path = Path(self._normalize_path(".")) + for path in self._selected: + file_path = Path(path.file_path) + if file_path != default_path: + selected_paths.add(file_path) + return selected_paths + def _get_composite_path(self, file_or_dir: str) -> _CompositePath: head, tail = os.path.split(file_or_dir) file_name, *other = tail.split("::")