33import traceback
44from functools import wraps
55from pathlib import Path
6- from typing import Any , Callable , Optional , TypeVar , Union
6+ from typing import Any , Callable , Coroutine , NoReturn , Optional , TypeVar , Union
77
88import pendulum
99import typer
2929from .client import initialize_client_sync
3030
3131YamlFileVar = TypeVar ("YamlFileVar" , bound = YamlFile )
32+ T = TypeVar ("T" )
3233
3334
3435def init_logging (debug : bool = False ) -> None :
@@ -42,7 +43,7 @@ def init_logging(debug: bool = False) -> None:
4243 logging .getLogger ("infrahubctl" )
4344
4445
45- def handle_exception (exc : Exception , console : Console , exit_code : int ):
46+ def handle_exception (exc : Exception , console : Console , exit_code : int ) -> NoReturn :
4647 """Handle exeception in a different fashion based on its type."""
4748 if isinstance (exc , Exit ):
4849 raise typer .Exit (code = exc .exit_code )
@@ -67,16 +68,18 @@ def handle_exception(exc: Exception, console: Console, exit_code: int):
6768 raise typer .Exit (code = exit_code )
6869
6970
70- def catch_exception (console : Optional [Console ] = None , exit_code : int = 1 ):
71+ def catch_exception (
72+ console : Optional [Console ] = None , exit_code : int = 1
73+ ) -> Callable [[Callable [..., T ]], Callable [..., Union [T , Coroutine [Any , Any , T ], NoReturn ]]]:
7174 """Decorator to handle exception for commands."""
7275 if not console :
7376 console = Console ()
7477
75- def decorator (func : Callable ) :
78+ def decorator (func : Callable [..., T ]) -> Callable [..., Union [ T , Coroutine [ Any , Any , T ], NoReturn ]] :
7679 if asyncio .iscoroutinefunction (func ):
7780
7881 @wraps (func )
79- async def async_wrapper (* args : Any , ** kwargs : Any ):
82+ async def async_wrapper (* args : Any , ** kwargs : Any ) -> Union [ T , NoReturn ] :
8083 try :
8184 return await func (* args , ** kwargs )
8285 except (Error , Exception ) as exc : # pylint: disable=broad-exception-caught
@@ -85,7 +88,7 @@ async def async_wrapper(*args: Any, **kwargs: Any):
8588 return async_wrapper
8689
8790 @wraps (func )
88- def wrapper (* args : Any , ** kwargs : Any ):
91+ def wrapper (* args : Any , ** kwargs : Any ) -> Union [ T , NoReturn ] :
8992 try :
9093 return func (* args , ** kwargs )
9194 except (Error , Exception ) as exc : # pylint: disable=broad-exception-caught
@@ -116,8 +119,10 @@ def execute_graphql_query(
116119 )
117120
118121 if debug :
119- message = ("-" * 40 , f"Response for GraphQL Query { query } " , response , "-" * 40 )
120- console .print ("\n " .join (message ))
122+ console .print ("-" * 40 )
123+ console .print (f"Response for GraphQL Query { query } " )
124+ console .print (response )
125+ console .print ("-" * 40 )
121126
122127 return response
123128
0 commit comments