Coverage for src/cell_abm_pipeline/utilities/prefect_plugin.py: 0%

1from mypy.plugin import Plugin 

2from mypy.types import CallableType 

3 

4 

5class CustomPlugin(Plugin): 

6 def get_method_hook(self, fullname: str): 

7 if fullname == "prefect.tasks.Task.__call__": 

8 return update_prefect_task_call_return 

9 

10 return None 

11 

12 def get_method_signature_hook(self, fullname: str): 

13 if fullname == "prefect.tasks.Task.submit": 

14 return update_prefect_task_submit_signature 

15 

16 return None 

17 

18 

19def update_prefect_task_call_return(ctx): 

20 # Replace default None return of Task with return type of underlying method. 

21 return ctx.type.args[1] 

22 

23 

24def update_prefect_task_submit_signature(ctx): 

25 # Drop the return_state argument for submit task calls. 

26 if "return_state" in ctx.default_signature.arg_names: 

27 index = ctx.default_signature.arg_names.index("return_state") 

28 return CallableType( 

29 ctx.default_signature.arg_types[:index] + ctx.default_signature.arg_types[index + 1 :], 

30 ctx.default_signature.arg_kinds[:index] + ctx.default_signature.arg_kinds[index + 1 :], 

31 ctx.default_signature.arg_names[:index] + ctx.default_signature.arg_names[index + 1 :], 

32 ctx.default_signature.ret_type, 

33 ctx.default_signature.fallback, 

34 ) 

35 

36 return ctx.default_signature 

37 

38 

39def plugin(version: str): 

40 return CustomPlugin