make fmt & make mypy

Signed-off-by: Samhita Alla <aallasamhita@gmail.com>

Author: samhita-alla

examples/shell/01_basic.py

@@ -8,12 +8,12 @@
 
 import sys
 import tempfile
+from pathlib import Path
 
 import flyte
+from flyte._image import PythonWheels
 from flyte.extras import shell
 from flyte.io import File
-from flyte._image import PythonWheels
-from pathlib import Path
 
 # Wrap `head` — emits the first N lines of an input file.
 head_task = shell.create(
@@ -60,8 +60,6 @@ async def take_first_lines(src: File, n: int) -> File:
         f.write("\n".join(f"line {i}" for i in range(1, 21)))
         path = f.name
 
-    run = flyte.with_runcontext(mode=mode).run(
-        take_first_lines, File.from_local_sync(path), 5
-    )
+    run = flyte.with_runcontext(mode=mode).run(take_first_lines, File.from_local_sync(path), 5)
     print(run.url if mode == "remote" else run)
     print(f"Output: {run.outputs()}")

examples/shell/02_lists_and_globs.py

@@ -77,9 +77,7 @@ async def split_concatenated(parts: list[File], chunk_lines: int) -> list[File]:
 
     parts: list[File] = []
     for i in range(3):
-        with tempfile.NamedTemporaryFile(
-            mode="w", suffix=f"_p{i}.txt", delete=False
-        ) as f:
+        with tempfile.NamedTemporaryFile(mode="w", suffix=f"_p{i}.txt", delete=False) as f:
             f.write("\n".join(f"part-{i}/line-{j}" for j in range(10)))
             parts.append(File.from_local_sync(f.name))
 

examples/shell/04_dict_inputs.py

@@ -46,9 +46,7 @@
     image="debian:12-slim",
     inputs={"opts": dict[str, str]},
     outputs={"argv": Stdout(type=str)},
-    flag_aliases={
-        "opts": ""
-    },  # pairs mode is the default; no per-key prefix; keys already include `--`
+    flag_aliases={"opts": ""},  # pairs mode is the default; no per-key prefix; keys already include `--`
     script=r"""
         echo {flags.opts}
     """,
@@ -117,7 +115,7 @@ async def dict_demo() -> tuple[str, str, str, str]:
         opts={
             "--memory": "4G",
             "--threads": "8",
-            "--label": "it's a test",      # single quote — safe via positional args
+            "--label": "it's a test",  # single quote — safe via positional args
         }
     )
     equals_argv = await echo_equals(

examples/shell/05_bool_and_optional.py

@@ -71,9 +71,7 @@ async def bool_demo() -> tuple[str, str, str]:
     # All flags on.
     a = await report(title="full", verbose=True, case_insensitive=True, threads=8)
     # Only verbose; threads omitted (None).
-    b = await report(
-        title="partial", verbose=True, case_insensitive=False, threads=None
-    )
+    b = await report(title="partial", verbose=True, case_insensitive=False, threads=None)
     # All flags off.
     c = await report(title="bare", verbose=False, case_insensitive=False, threads=None)
     return a, b, c

examples/shell/08_dirs.py

@@ -79,8 +79,6 @@ async def summarize_files(src: Dir) -> Dir:
     for i, body in enumerate(["one\ntwo\nthree\n", "alpha\nbeta\n", "single\n"]):
         (tmp / f"file_{i}.txt").write_text(body)
 
-    run = flyte.with_runcontext(mode=mode).run(
-        summarize_files, Dir.from_local_sync(str(tmp))
-    )
+    run = flyte.with_runcontext(mode=mode).run(summarize_files, Dir.from_local_sync(str(tmp)))
     print(run.url if mode == "remote" else run)
     print(f"Output: {run.outputs()}")

examples/shell/10_resources_and_image.py

@@ -38,9 +38,7 @@
 # A layered Image — base + extra apt package. shell.create builds this on
 # first call via flyte.build (using the configured builder) and passes the
 # resolved URI down to ContainerTask.
-image_with_jq = flyte.Image.from_debian_base(install_flyte=False).with_apt_packages(
-    "jq"
-)
+image_with_jq = flyte.Image.from_debian_base(install_flyte=False).with_apt_packages("jq")
 
 
 pretty_print_json = shell.create(

examples/shell/11_list_flag_modes.py

@@ -74,9 +74,7 @@ async def list_mode_demo(parts: list[File]) -> tuple[str, str, str]:
 
     parts: list[File] = []
     for i in range(3):
-        with tempfile.NamedTemporaryFile(
-            mode="w", suffix=f"_{i}.txt", delete=False
-        ) as f:
+        with tempfile.NamedTemporaryFile(mode="w", suffix=f"_{i}.txt", delete=False) as f:
             f.write(f"part-{i}\n")
             parts.append(File.from_local_sync(f.name))
 

src/flyte/extras/shell/__init__.py

@@ -57,36 +57,36 @@
 from __future__ import annotations
 
 from ._render import _DICT_SEP, _render_command
-from ._runtime import _Shell, _read_process_result, create
+from ._runtime import _read_process_result, _Shell, create
 from ._types import (
     DictMode,
     FlagSpec,
     Glob,
     OutputSpec,
     Stderr,
     Stdout,
-    listMode,
     _classify_input,
     _is_list_of,
     _is_optional,
     _validate_outputs,
+    listMode,
 )
 
 __all__ = [
+    "_DICT_SEP",
     "DictMode",
     "FlagSpec",
     "Glob",
     "OutputSpec",
     "Stderr",
     "Stdout",
-    "create",
-    "listMode",
-    "_DICT_SEP",
     "_Shell",
     "_classify_input",
     "_is_list_of",
     "_is_optional",
     "_read_process_result",
     "_render_command",
     "_validate_outputs",
+    "create",
+    "listMode",
 ]

src/flyte/extras/shell/_render.py

@@ -7,7 +7,6 @@
 
 from ._types import FlagSpec, Stderr, Stdout, _classify_input
 
-
 _PLACEHOLDER_RE = re.compile(r"\{(inputs|flags|outputs)\.([a-zA-Z_][a-zA-Z0-9_]*)\}")
 _DICT_SEP = "\x1e"
 
@@ -57,9 +56,7 @@ def render_input_ref(name: str) -> str:
         kind = kinds.get(name)
 
         if kind is None:
-            raise KeyError(
-                f"{{inputs.{name}}} used in script but {name!r} is not declared in inputs."
-            )
+            raise KeyError(f"{{inputs.{name}}} used in script but {name!r} is not declared in inputs.")
         if kind in ("file", "dir"):
             return str(input_data_dir / name)
         if kind == "list_file":
@@ -108,10 +105,7 @@ def replace(match: re.Match) -> str:
             return render_flag_ref(name)
         if ns == "outputs":
             if name not in outputs:
-                raise KeyError(
-                    f"{{outputs.{name}}} references an unknown output. "
-                    f"Declared outputs: {list(outputs)}"
-                )
+                raise KeyError(f"{{outputs.{name}}} references an unknown output. Declared outputs: {list(outputs)}")
             spec = outputs[name]
             if isinstance(spec, (Stdout, Stderr)):
                 raise KeyError(
@@ -151,11 +145,7 @@ def _emit_flag_setter(
 
     if kind == "bool":
         val_var = alloc_slot(name)
-        return (
-            f'if [ "${{{val_var}}}" = "true" ]; then '
-            f"{flag_var}={shlex.quote(flag)}; "
-            f'else {flag_var}=""; fi'
-        )
+        return f'if [ "${{{val_var}}}" = "true" ]; then {flag_var}={shlex.quote(flag)}; else {flag_var}=""; fi'
     if kind == "scalar":
         val_var = alloc_slot(name)
         return (

src/flyte/extras/shell/_runtime.py

@@ -11,15 +11,15 @@
 
 from ._render import _DICT_SEP, _render_command
 from ._types import (
+    _SCALAR_TYPES,
     FlagSpec,
     Glob,
     Stderr,
     Stdout,
-    _ProcessResult,
-    _SCALAR_TYPES,
     _classify_input,
     _is_dict_str_str,
     _is_optional,
+    _ProcessResult,
     _validate_outputs,
     listMode,
 )
@@ -51,8 +51,8 @@ class _Shell:
     _env: Optional["flyte.TaskEnvironment"] = field(default=None, repr=False, compare=False)
     _resolved_image_uri: Optional[str] = field(default=None, repr=False, compare=False)
 
-    def _container_inputs(self) -> dict[str, Type]:
-        wired: dict[str, Type] = {}
+    def _container_inputs(self) -> dict[str, Any]:
+        wired: dict[str, Any] = {}
         for name, tp in self.inputs.items():
             is_opt, inner = _is_optional(tp)
             if _is_dict_str_str(inner):
@@ -81,10 +81,7 @@ def _build_command(self) -> list[str]:
             script=self.script,
             inputs=self.inputs,
             outputs=self.outputs,
-            flag_specs={
-                name: FlagSpec.coerce(name, self.flag_aliases.get(name))
-                for name in self.inputs
-            },
+            flag_specs={name: FlagSpec.coerce(name, self.flag_aliases.get(name)) for name in self.inputs},
             input_data_dir=self.input_data_dir,
             output_data_dir=self.output_data_dir,
         )
@@ -100,17 +97,13 @@ def _build_command(self) -> list[str]:
         mkdirs = [
             f"mkdir -p {shlex.quote(str(self.output_data_dir / name))}"
             for name, spec in self.outputs.items()
-            if isinstance(spec, Glob)
-            or (isinstance(spec, type) and issubclass(spec, Dir))
+            if isinstance(spec, Glob) or (isinstance(spec, type) and issubclass(spec, Dir))
         ]
         mkdir_preamble = "; ".join(mkdirs) + ";" if mkdirs else ""
 
         debug_preamble = ""
         if self.debug:
-            debug_preamble = (
-                'echo "--- shell task: rendered script ---" >&2; '
-                f"cat <<'_EOF_' >&2\n{body}\n_EOF_\n"
-            )
+            debug_preamble = f"echo \"--- shell task: rendered script ---\" >&2; cat <<'_EOF_' >&2\n{body}\n_EOF_\n"
 
         wrapped = (
             f"{mkdir_preamble} "
@@ -187,9 +180,7 @@ async def _unpack_outputs(self, raw: Any) -> Any:
             if isinstance(spec, Glob) and isinstance(value, Dir):
                 local = await value.download() if hasattr(value, "download") else value.path
                 matched = sorted(pathlib.Path(str(local)).glob(spec.pattern))
-                unpacked.append(
-                    [await File.from_local(str(p)) for p in matched if p.is_file()]
-                )
+                unpacked.append([await File.from_local(str(p)) for p in matched if p.is_file()])
             else:
                 unpacked.append(value)
         return unpacked[0] if single else tuple(unpacked)
@@ -223,20 +214,15 @@ async def _prepare_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
                 for k, v in value.items():
                     if _DICT_SEP in k or _DICT_SEP in v:
                         raise ValueError(
-                            f"dict input {name!r}: keys/values cannot contain the "
-                            f"record-separator byte (\\x1e)."
+                            f"dict input {name!r}: keys/values cannot contain the record-separator byte (\\x1e)."
                         )
                     parts.append(k)
                     parts.append(v)
                 out[name] = _DICT_SEP.join(parts)
                 continue
 
             if is_opt and kind in ("scalar", "bool"):
-                out[name] = (
-                    "true"
-                    if (kind == "bool" and value)
-                    else "false" if kind == "bool" else str(value)
-                )
+                out[name] = "true" if (kind == "bool" and value) else "false" if kind == "bool" else str(value)
                 continue
 
             out[name] = value
@@ -310,9 +296,7 @@ def create(
     inputs: Optional[dict[str, Type]] = None,
     outputs: Optional[dict[str, Any]] = None,
     script: str,
-    flag_aliases: Optional[
-        dict[str, Union[str, Tuple[str, listMode], FlagSpec]]
-    ] = None,
+    flag_aliases: Optional[dict[str, Union[str, Tuple[str, listMode], FlagSpec]]] = None,
     shell: str = "/bin/bash",
     debug: bool = False,
     resources: Optional[flyte.Resources] = None,
@@ -484,16 +468,11 @@ async def pipeline(a: File, b: list[File]) -> list[File]:
     coerced_aliases: dict[str, FlagSpec] = {}
     for n, alias in (flag_aliases or {}).items():
         if n not in inputs:
-            raise KeyError(
-                f"flag_aliases references {n!r} which is not declared in inputs."
-            )
+            raise KeyError(f"flag_aliases references {n!r} which is not declared in inputs.")
         coerced_aliases[n] = FlagSpec.coerce(n, alias)
 
     if not isinstance(image, (str, flyte.Image)):
-        raise TypeError(
-            f"image must be a URI string or a flyte.Image, got "
-            f"{type(image).__name__}."
-        )
+        raise TypeError(f"image must be a URI string or a flyte.Image, got {type(image).__name__}.")
 
     return _Shell(
         name=name,