1-
21from __future__ import annotations
32
43import logging
@@ -52,7 +51,7 @@ def __init__(
5251 self .version = version
5352 self .tag = tag
5453
55- super ().__init__ (** kwargs ) # type: ignore [reportUnknownMemberType]
54+ super ().__init__ (** kwargs ) # type: ignore [reportUnknownMemberType]
5655
5756 def execute (self , context : Any ) -> None :
5857 """Operator execute method.
@@ -95,8 +94,7 @@ def __init__(
9594 self .stage = stage
9695 self .model_name = model_name
9796
98- super ().__init__ (** kwargs ) # type: ignore [reportUnknownMemberType]
99-
97+ super ().__init__ (** kwargs ) # type: ignore [reportUnknownMemberType]
10098
10199 def execute (self , context : Any ) -> list [SlimModelVersion ]:
102100 """Connect to Mlflow hook and dump messages to tmp file."""
@@ -131,7 +129,7 @@ def __init__(
131129 self .tracking_uri = tracking_uri
132130 self .upstream_task_id = upstream_task_id
133131
134- super ().__init__ (** kwargs )# type: ignore [reportUnknownMemberType]
132+ super ().__init__ (** kwargs ) # type: ignore [reportUnknownMemberType]
135133
136134 def execute (self , context : Any ) -> Any :
137135 """Operator execute method.
@@ -147,13 +145,13 @@ def execute(self, context: Any) -> Any:
147145 key = "return_value" , task_ids = self .upstream_task_id
148146 )
149147 if isinstance (xcom_val , list ):
150- _first_item = xcom_val .pop () # type: ignore [unknownTypeIssue]
151- run_id = _first_item ["run_id" ]# type: ignore [unknownTypeIssue]
148+ _first_item = xcom_val .pop () # type: ignore [unknownTypeIssue]
149+ run_id = _first_item ["run_id" ] # type: ignore [unknownTypeIssue]
152150 logging .warning (
153151 "Multiple items returned from upstream task. Defaulting to first item"
154152 )
155153 elif isinstance (xcom_val , dict ):
156- run_id = xcom_val ["run_id" ]# type: ignore [unknownTypeIssue]
154+ run_id = xcom_val ["run_id" ] # type: ignore [unknownTypeIssue]
157155 else :
158156 logging .error (
159157 "Operator has no implementation for this type: %s" , f"{ type (xcom_val )} ."
@@ -164,8 +162,7 @@ def execute(self, context: Any) -> Any:
164162 raise ValueError (f"{ run_id } must be an instance of 'str'" )
165163 data = client_hook .get_run (run_id )
166164 logging .info ("%s" , f"{ data = } " )
167- return data .data .metrics # type: ignore [unknownTypeIssue]
168-
165+ return data .data .metrics # type: ignore [unknownTypeIssue]
169166
170167
171168class GetLatestModelVersion (BaseOperator ):
@@ -192,7 +189,7 @@ def __init__(
192189 self .tracking_uri = tracking_uri
193190 self .model_name = model_name
194191
195- super ().__init__ (** kwargs )# type: ignore [reportUnknownMemberType]
192+ super ().__init__ (** kwargs ) # type: ignore [reportUnknownMemberType]
196193
197194 def execute (self , context : Any ) -> SlimModelVersion :
198195 """Operator execute method.
0 commit comments