diff --git a/models/extractor.py b/models/extractor.py new file mode 100644 index 0000000..760160c --- /dev/null +++ b/models/extractor.py @@ -0,0 +1,31 @@ +from typing import Any +from .context import Context +from .itm import ExtractModel +from .executor import Response + + +class Extractor: + def __init__(self, context: Context): + self.context = context + + def extract(self, configs: list[ExtractModel], response: Response): + for config in configs: + value = self._get_value(response, config) + self.context.set_extracted(config.name, value) + + def _get_value(self, response: Response, config: ExtractModel) -> Any: + if config.source == "status": + return response.status + if config.source == "header": + return response.headers.get(config.json_path) + return self._json_path(response.body, config.json_path) + + def _json_path(self, data: dict, path: str) -> Any: + parts = path.lstrip("$.").split(".") + current = data + for part in parts: + if isinstance(current, dict): + current = current.get(part) + else: + return None + return current