diff --git a/fireworks/core/firework.py b/fireworks/core/firework.py index 8f9f0d1eb..0bc9d46fc 100644 --- a/fireworks/core/firework.py +++ b/fireworks/core/firework.py @@ -14,7 +14,7 @@ from collections import OrderedDict, defaultdict from copy import deepcopy from datetime import datetime -from typing import Any, Dict, Iterable, List, Optional, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union from monty.io import reverse_readline, zopen from monty.os.path import zpath diff --git a/fireworks/scripts/lpad_run.py b/fireworks/scripts/lpad_run.py index 15c9e627a..e8d32f7e4 100644 --- a/fireworks/scripts/lpad_run.py +++ b/fireworks/scripts/lpad_run.py @@ -11,7 +11,7 @@ import time import traceback from argparse import ArgumentParser, ArgumentTypeError, Namespace -from typing import Any, Callable, Sequence +from typing import Any, Callable, List, Sequence, Union import ruamel.yaml as yaml from pymongo import ASCENDING, DESCENDING @@ -317,7 +317,7 @@ def get_fw_ids_helper(lp, args, count_only=None): return ids -def get_fws_helper(lp: LaunchPad, ids: Sequence[int], args: Namespace) -> List[Any]: +def get_fws_helper(lp: LaunchPad, ids: List[int], args: Namespace) -> Union[int, List[int], List[List[int]]]: """Get fws from ids in a representation according to args.display_format.""" fws = [] if args.display_format == "ids":