added type annotation

mitm_tooling/utilities/python_utils.py

@@ -90,18 +90,33 @@ def inner_list_concat(d1: dict[K, list[Any]], d2: dict[K, list[Any]]) -> dict[K,
     return res
 
 
-def pick_from_mapping(d: Mapping[K, T], keys: Sequence[K]) -> list[tuple[K, T]]:
-    return [(k, d[k]) for k in keys]
+@overload
+def pick_from_mapping(d: Mapping[K, T], keys: Sequence[K], *, flatten: bool = True) -> list[T]: ...
 
 
 @overload
-def unique(*its: Iterable[T]) -> list[T]:
-    ...
+def pick_from_mapping(d: Mapping[K, T], keys: Sequence[K], *, flatten: bool = False) -> list[tuple[K, T]]: ...
+
+
+@overload
+def pick_from_mapping(d: Mapping[K, T], keys: Sequence[K], *, flatten: bool) -> list[T] | list[tuple[K, T]]: ...
+
+# TODO maybe fix type annotations?
+def pick_from_mapping(d: Mapping[K, T], keys: Sequence[K], *, flatten: bool = False) -> list[T] | list[tuple[K, T]]:
+    if flatten:
+        return [d[k] for k in keys]
+    else:
+        return [(k, d[k]) for k in keys]
 
 
 Ts = TypeVarTuple('Ts')
 
 
+@overload
+def unique(*its: Iterable[T]) -> list[T]:
+    ...
+
+
 def unique(*its: Iterable[Union[*Ts]]) -> list[Union[Unpack[Ts]]]:
     return list(set(itertools.chain(*its)))