Coverage for src/trapi_predict_kit/decorators.py: 89%
15 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-25 21:14 +0200
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-25 21:14 +0200
1import functools
2from typing import Any, Callable, Dict, List, Optional
4from reasoner_pydantic import MetaEdge, MetaNode
6from trapi_predict_kit.types import PredictOptions
9def trapi_predict(
10 path: str,
11 edges: List[MetaEdge],
12 nodes: Dict[str, MetaNode],
13 name: Optional[str] = None,
14 description: Optional[str] = "",
15 default_input: Optional[str] = "drugbank:DB00394",
16 default_model: Optional[str] = "openpredict_baseline",
17) -> Callable:
18 """A decorator to indicate a function is a function to generate prediction that can be integrated to TRAPI.
19 The function needs to take an input_id and returns a list of predicted entities related to the input entity
20 """
21 if not name: 21 ↛ 22line 21 didn't jump to line 22, because the condition on line 21 was never true
22 name = path
24 def decorator(func: Callable) -> Any:
25 @functools.wraps(func)
26 def wrapper(input_id: str, options: Optional[PredictOptions] = None) -> Callable:
27 options = PredictOptions.parse_obj(options) if options else PredictOptions()
28 return func(input_id, options)
30 wrapper._trapi_predict = {
31 "edges": edges,
32 "nodes": nodes,
33 "path": path,
34 "name": name,
35 "description": description,
36 "default_input": default_input,
37 "default_model": default_model,
38 }
40 return wrapper
42 return decorator