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

1import functools 

2from typing import Any, Callable, Dict, List, Optional 

3 

4from reasoner_pydantic import MetaEdge, MetaNode 

5 

6from trapi_predict_kit.types import PredictOptions 

7 

8 

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 

23 

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) 

29 

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 } 

39 

40 return wrapper 

41 

42 return decorator