11import enum
22import json
3+ from typing import Any , Dict , Optional , Union
34
45
56class CommandState (enum .Enum ):
@@ -9,24 +10,29 @@ class CommandState(enum.Enum):
910
1011class CommandRequest :
1112 @classmethod
12- def unmarshal_json (cls , data ) :
13+ def unmarshal_json (cls , data : Union [ str , bytes ]) -> "CommandRequest" :
1314 req = json .loads (data )
1415 return cls (id_ = req ["id" ], name = req ["name" ], args = req .get ("arguments" ))
1516
16- def __init__ (self , id_ , name , args = None ):
17+ def __init__ (self , id_ : str , name : str , args : Optional [ Dict [ str , Any ]] = None ):
1718 self .id = id_
1819 self .name = name
1920
2021 if args is None :
2122 args = {}
2223 self .args = args
2324
24- def new_response (self , * args , ** kwargs ):
25+ def new_response (self , * args , ** kwargs ) -> "CommandResponse" :
2526 return CommandResponse (self .id , * args , ** kwargs )
2627
2728
2829class CommandResponse :
29- def __init__ (self , id_ , state , payload = None ):
30+ def __init__ (
31+ self ,
32+ id_ : str ,
33+ state : Union [str , CommandState ],
34+ payload : Optional [Union [Dict [str , Any ], str ]] = None ,
35+ ) -> None :
3036 self .id = id_
3137
3238 if not isinstance (state , CommandState ):
@@ -37,8 +43,8 @@ def __init__(self, id_, state, payload=None):
3743 payload = {"message" : payload }
3844 self .payload = payload
3945
40- def json (self ):
41- json_object = {"id" : self .id , "state" : self .state .value }
46+ def json (self ) -> Dict [ str , Any ] :
47+ json_object : Dict [ str , Any ] = {"id" : self .id , "state" : self .state .value }
4248 if self .payload is not None :
4349 json_object ["payload" ] = self .payload
4450
0 commit comments