31
31
from airflow .utils .context import Context
32
32
from airflow .utils .edgemodifier import EdgeModifier
33
33
from airflow .utils .session import NEW_SESSION , provide_session
34
- from airflow .utils .types import NOTSET
34
+ from airflow .utils .types import NOTSET , ArgNotSet
35
35
36
36
if TYPE_CHECKING :
37
37
from airflow .models .dag import DAG
@@ -322,7 +322,7 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None:
322
322
def resolve (self , context : Context , session : Session = NEW_SESSION ) -> Any :
323
323
task_id = self .operator .task_id
324
324
result = context ["ti" ].xcom_pull (task_ids = task_id , key = str (self .key ), default = NOTSET , session = session )
325
- if result is not NOTSET :
325
+ if not isinstance ( result , ArgNotSet ) :
326
326
return result
327
327
if self .key == XCOM_RETURN_KEY :
328
328
return None
@@ -437,7 +437,7 @@ def __getitem__(self, index: Any) -> Any:
437
437
438
438
def __len__ (self ) -> int :
439
439
lengths = (len (v ) for v in self .values )
440
- if self .fillvalue is NOTSET :
440
+ if isinstance ( self .fillvalue , ArgNotSet ) :
441
441
return min (lengths )
442
442
return max (lengths )
443
443
@@ -460,13 +460,13 @@ def __repr__(self) -> str:
460
460
args_iter = iter (self .args )
461
461
first = repr (next (args_iter ))
462
462
rest = ", " .join (repr (arg ) for arg in args_iter )
463
- if self .fillvalue is NOTSET :
463
+ if isinstance ( self .fillvalue , ArgNotSet ) :
464
464
return f"{ first } .zip({ rest } )"
465
465
return f"{ first } .zip({ rest } , fillvalue={ self .fillvalue !r} )"
466
466
467
467
def _serialize (self ) -> dict [str , Any ]:
468
468
args = [serialize_xcom_arg (arg ) for arg in self .args ]
469
- if self .fillvalue is NOTSET :
469
+ if isinstance ( self .fillvalue , ArgNotSet ) :
470
470
return {"args" : args }
471
471
return {"args" : args , "fillvalue" : self .fillvalue }
472
472
@@ -486,7 +486,7 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None:
486
486
ready_lengths = [length for length in all_lengths if length is not None ]
487
487
if len (ready_lengths ) != len (self .args ):
488
488
return None # If any of the referenced XComs is not ready, we are not ready either.
489
- if self .fillvalue is NOTSET :
489
+ if isinstance ( self .fillvalue , ArgNotSet ) :
490
490
return min (ready_lengths )
491
491
return max (ready_lengths )
492
492
0 commit comments