Fix some imports

materialsproject · Jan 28, 2022 · d1d6bb9 · d1d6bb9
1 parent 65eafcd
commit d1d6bb9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git 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
@@ -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":