1515import logging
1616from abc import ABCMeta , abstractmethod
1717from collections import defaultdict
18+ from typing import Dict , Set
1819
1920from slugid import nice as slugid
2021
@@ -365,11 +366,11 @@ def replace_tasks(
365366
366367
367368def get_subgraph (
368- target_task_graph ,
369- removed_tasks ,
370- replaced_tasks ,
371- label_to_taskid ,
372- decision_task_id ,
369+ target_task_graph : TaskGraph ,
370+ removed_tasks : Set [ str ] ,
371+ replaced_tasks : Set [ str ] ,
372+ label_to_taskid : Dict [ str , str ] ,
373+ decision_task_id : str ,
373374):
374375 """
375376 Return the subgraph of target_task_graph consisting only of
@@ -399,7 +400,9 @@ def get_subgraph(
399400 for label in sorted (
400401 target_task_graph .graph .nodes - removed_tasks - set (label_to_taskid )
401402 ):
402- label_to_taskid [label ] = slugid ()
403+ task_id = slugid ()
404+ assert isinstance (task_id , str )
405+ label_to_taskid [label ] = task_id
403406
404407 # resolve labels to taskIds and populate task['dependencies']
405408 tasks_by_taskid = {}
@@ -424,14 +427,15 @@ def get_subgraph(
424427 }
425428 )
426429
430+ assert task .task_id
427431 task .task = resolve_task_references (
428432 task .label ,
429433 task .task ,
430434 task_id = task .task_id ,
431435 decision_task_id = decision_task_id ,
432436 dependencies = named_task_dependencies ,
433437 )
434- deps = task .task .setdefault ("dependencies" , []) # type: ignore
438+ deps = task .task .setdefault ("dependencies" , [])
435439 deps .extend (sorted (named_task_dependencies .values ()))
436440 tasks_by_taskid [task .task_id ] = task
437441
0 commit comments