88
99import boto3 as boto
1010import httpx
11+ from typing import List
1112
1213from sync .clients .sync import get_default_client
1314from sync .models import Platform , PredictionError , Response
1415
1516logger = logging .getLogger (__name__ )
1617
1718
18- def get_products () -> Response [list [str ]]:
19+ def get_products () -> Response [List [str ]]:
1920 """Get supported platforms
2021
2122 :return: list of platform names
@@ -43,7 +44,8 @@ def generate_prediction(
4344 """
4445 response = create_prediction (platform , cluster_report , eventlog_url )
4546
46- if prediction_id := response .result :
47+ prediction_id = response .result
48+ if prediction_id :
4749 return wait_for_prediction (prediction_id , preference )
4850
4951 return response
@@ -60,8 +62,8 @@ def wait_for_prediction(prediction_id: str, preference: str = None) -> Response[
6062 :rtype: Response[dict]
6163 """
6264 response = wait_for_final_prediction_status (prediction_id )
63-
64- if result := response . result :
65+ result = response . result
66+ if result :
6567 if result == "SUCCESS" :
6668 return get_prediction (prediction_id , preference )
6769
@@ -84,7 +86,8 @@ def get_prediction(prediction_id: str, preference: str = None) -> Response[dict]
8486 prediction_id , {"preference" : preference } if preference else None
8587 )
8688
87- if result := response .get ("result" ):
89+ result = response .get ("result" )
90+ if result :
8891 return Response (result = result )
8992
9093 return Response (** response )
@@ -100,15 +103,16 @@ def get_status(prediction_id: str) -> Response[str]:
100103 """
101104 response = get_default_client ().get_prediction_status (prediction_id )
102105
103- if result := response .get ("result" ):
106+ result = response .get ("result" )
107+ if result :
104108 return Response (result = result ["status" ])
105109
106110 return Response (** response )
107111
108112
109113def get_predictions (
110114 product : str = None , project_id : str = None , preference : str = None
111- ) -> Response [list [dict ]]:
115+ ) -> Response [List [dict ]]:
112116 """Get predictions
113117
114118 :param product: platform to filter by, e.g. "aws-emr", defaults to None
@@ -144,8 +148,10 @@ def wait_for_final_prediction_status(prediction_id: str) -> Response[str]:
144148 :return: prediction status, e.g. "SUCCESS"
145149 :rtype: Response[str]
146150 """
147- while response := get_default_client ().get_prediction_status (prediction_id ):
148- if result := response .get ("result" ):
151+ response = get_default_client ().get_prediction_status (prediction_id )
152+ while response :
153+ result = response .get ("result" )
154+ if result :
149155 if result ["status" ] in ("SUCCESS" , "FAILURE" ):
150156 return Response (result = result ["status" ])
151157 else :
@@ -154,6 +160,8 @@ def wait_for_final_prediction_status(prediction_id: str) -> Response[str]:
154160 logger .info ("Waiting for prediction" )
155161 sleep (10 )
156162
163+ response = get_default_client ().get_prediction_status (prediction_id )
164+
157165 return Response (error = PredictionError (message = "Failed to get prediction status" ))
158166
159167
@@ -221,16 +229,16 @@ def create_prediction(
221229 :return: prediction ID
222230 :rtype: Response[str]
223231 """
224- match urlparse (eventlog_url ).scheme :
225- case "s3" :
226- response = generate_presigned_url (eventlog_url )
227- if response .error :
228- return response
229- eventlog_http_url = response .result
230- case "http" | "https" :
231- eventlog_http_url = eventlog_url
232- case _ :
233- return Response (error = PredictionError (message = "Unsupported event log URL scheme" ))
232+ scheme = urlparse (eventlog_url ).scheme
233+ if scheme == "s3" :
234+ response = generate_presigned_url (eventlog_url )
235+ if response .error :
236+ return response
237+ eventlog_http_url = response .result
238+ elif scheme in { "http" , "https" } :
239+ eventlog_http_url = eventlog_url
240+ else :
241+ return Response (error = PredictionError (message = "Unsupported event log URL scheme" ))
234242
235243 response = get_default_client ().create_prediction (
236244 {
0 commit comments