From ef61b51368373df8b527f2c4737cb7a07c8ca0bd Mon Sep 17 00:00:00 2001
From: Ketan Umare <ketan.umare@gmail.com>
Date: Wed, 10 Nov 2021 16:59:52 -0800
Subject: [PATCH 1/6] [wip] Shell Task

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>
---
 flytekit/extras/tasks/__init__.py             |   0
 flytekit/extras/tasks/shell.py                | 208 ++++++++++++++++++
 flytekit/types/directory/types.py             |   9 +-
 flytekit/types/file/file.py                   |   9 +-
 tests/flytekit/unit/extras/tasks/__init__.py  |   0
 .../flytekit/unit/extras/tasks/test_shell.py  |  76 +++++++
 .../unit/extras/tasks/testdata/test.csv       |   0
 7 files changed, 286 insertions(+), 16 deletions(-)
 create mode 100644 flytekit/extras/tasks/__init__.py
 create mode 100644 flytekit/extras/tasks/shell.py
 create mode 100644 tests/flytekit/unit/extras/tasks/__init__.py
 create mode 100644 tests/flytekit/unit/extras/tasks/test_shell.py
 create mode 100644 tests/flytekit/unit/extras/tasks/testdata/test.csv

diff --git a/flytekit/extras/tasks/__init__.py b/flytekit/extras/tasks/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
new file mode 100644
index 0000000000..d964173fe9
--- /dev/null
+++ b/flytekit/extras/tasks/shell.py
@@ -0,0 +1,208 @@
+import datetime
+import logging
+import os
+import re
+import subprocess
+import typing
+from dataclasses import dataclass
+
+from flytekit import PythonInstanceTask, ExecutionParameters
+from flytekit.core.interface import Interface
+from flytekit.core.task import TaskPlugins
+from flytekit.types.directory import FlyteDirectory
+from flytekit.types.file import FlyteFile
+
+
+@dataclass
+class OutputLocation:
+    """
+    Args:
+        var: str The name of the output variable
+        var_type: typing.Type The type of output variable
+        location: os.PathLike The location where this output variable will be written to or a regex that accepts input
+                  vars and generates the path. Of the form ``"{{ .inputs.v }}.tmp.md"``.
+                  This example for a given input v, at path `/tmp/abc.csv` will resolve to `/tmp/abc.csv.tmp.md`
+    """
+    var: str
+    var_type: typing.Type
+    location: typing.Union[os.PathLike, str]
+
+
+def _stringify(v: typing.Any) -> str:
+    """
+    Special cased return for the given value. Given the type returns the string version for the type.
+    Handles FlyteFile and FlyteDirectory specially. Downloads and returns the downloaded filepath
+    """
+    if isinstance(v, FlyteFile):
+        v.download()
+        return v.path
+    if isinstance(v, FlyteDirectory):
+        v.download()
+        return v.path
+    if isinstance(v, datetime.datetime):
+        return v.isoformat()
+    return str(v)
+
+
+def _interpolate(tmpl: str, regex: re.Pattern, validate_all_match: bool = True, **kwargs) -> str:
+    """
+    Substitutes all templates that match the input regex - ``r"({{\s*.inputs.(\w+)\s*}})", re.IGNORECASE`` with the
+    given inputs and returns the substituted string. The result is non destructive towards the given string.
+    """
+    modified = tmpl
+    matched = set()
+    for match in regex.finditer(tmpl):
+        expr = match.groups()[0]
+        var = match.groups()[1]
+        if var not in kwargs:
+            raise ValueError(f"Variable {var} in Query (part of {expr}) not found in inputs {kwargs.keys()}")
+        matched.add(var)
+        val = kwargs[var]
+        # str conversion should be deliberate, with right conversion for each type
+        modified = modified.replace(expr, _stringify(val))
+
+    if validate_all_match:
+        if len(matched) < len(kwargs.keys()):
+            diff = set(kwargs.keys()).difference(matched)
+            raise ValueError(f"Extra Inputs have no matches in query template - missing {diff}")
+    return modified
+
+
+def _dummy_task_func():
+    """
+    A Fake function to satisfy the inner PythonTask requirements
+    """
+    return None
+
+
+T = typing.TypeVar("T")
+
+
+class ShellTask(PythonInstanceTask[T]):
+    """
+    """
+
+    _INPUT_REGEX = re.compile(r"({{\s*.inputs.(\w+)\s*}})", re.IGNORECASE)
+    _OUTPUT_REGEX = re.compile(r"({{\s*.outputs.(\w+)\s*}})", re.IGNORECASE)
+
+    def __init__(
+            self,
+            name: str,
+            debug: bool = False,
+            script: typing.Optional[str] = None,
+            script_file: typing.Optional[os.PathLike] = None,
+            task_config: T = None,
+            inputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
+            output_locs: typing.Optional[typing.List[OutputLocation]] = None,
+            **kwargs,
+    ):
+        """
+        Args:
+            name: str Name of the Task. Should be unique in the project
+            debug: bool Print the generated script and other debugging information
+            script: The actual script specified as a string
+            script_file: A path to the file that contains the script (Only script or script_file) can be provided
+            task_config: T Configuration for the task, can be either a Pod (or coming soon, BatchJob) config
+            inputs: A Dictionary of input names to types
+            output_locs: A list of :ref:class:`OutputLocations`
+            **kwargs: Other arguments that can be passed to :ref:class:`PythonInstanceTask`
+        """
+        if script and script_file:
+            raise ValueError("Only either of script or script_file can be provided")
+        if not script and not script_file:
+            raise ValueError("Either a script or script_file is needed")
+        if script_file:
+            if not os.path.exists(script_file):
+                raise ValueError(f"FileNotFound: the specified Script file at path {script_file} cannot be loaded")
+
+        # Each instance of NotebookTask instantiates an underlying task with a dummy function that will only be used
+        # to run pre- and post- execute functions using the corresponding task plugin.
+        # We rename the function name here to ensure the generated task has a unique name and avoid duplicate task name
+        # errors.
+        # This seem like a hack. We should use a plugin_class that doesn't require a fake-function to make work.
+        plugin_class = TaskPlugins.find_pythontask_plugin(type(task_config))
+        self._config_task_instance = plugin_class(task_config=task_config, task_function=_dummy_task_func)
+        # Rename the internal task so that there are no conflicts at serialization time. Technically these internal
+        # tasks should not be serialized at all, but we don't currently have a mechanism for skipping Flyte entities
+        # at serialization time.
+        self._config_task_instance._name = f"_bash.{name}"
+        self._script = script
+        self._script_file = script_file
+        self._debug = debug
+        self._output_locs = output_locs if output_locs else []
+        outputs = self._validate_output_locs()
+        super().__init__(
+            name, task_config, task_type=self._config_task_instance.task_type,
+            interface=Interface(inputs=inputs, outputs=outputs), **kwargs
+        )
+
+    def _validate_output_locs(self) -> typing.Dict[str, typing.Type]:
+        outputs = {}
+        for v in self._output_locs:
+            if v is None:
+                raise ValueError("OutputLocation cannot be none")
+            if not isinstance(v, OutputLocation):
+                raise ValueError("Every output type should be an output location on the file-system")
+            if v.location is None:
+                raise ValueError(f"Output Location not provided for output var {v.var}")
+            if not issubclass(v.var_type, FlyteFile) and not issubclass(v.var_type, FlyteDirectory):
+                raise ValueError(f"Currently only outputs of type FlyteFile/FlyteDirectory and their derived types"
+                                 f" are supported")
+            outputs[v.var] = v.var_type
+        return outputs
+
+    @property
+    def script(self) -> typing.Optional[str]:
+        return self._script
+
+    @property
+    def script_file(self) -> typing.Optional[os.PathLike]:
+        return self._script_file
+
+    def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
+        return self._config_task_instance.pre_execute(user_params)
+
+    def execute(self, **kwargs) -> typing.Any:
+        """
+        Executes the given script by substituting the inputs and outputs and extracts the outputs from the filesystem
+        """
+        logging.info(f"Running shell script as type {self.task_type}")
+        if self.script_file:
+            with open(self.script_file) as f:
+                self._script = f.read()
+
+        outputs: typing.Dict[str, str] = {}
+        if self._output_locs:
+            for v in self._output_locs:
+                outputs[v.var] = _interpolate(v.location, self._INPUT_REGEX, validate_all_match=False, **kwargs)
+
+        gen_script = _interpolate(self._script, self._INPUT_REGEX, **kwargs)
+        # For outputs it is not necessary that all outputs are used in the script, some are implicit outputs
+        # for example gcc main.c will generate a.out automatically
+        gen_script = _interpolate(gen_script, self._OUTPUT_REGEX, validate_all_match=False, **outputs)
+        if self._debug:
+            print("\n==============================================\n")
+            print(gen_script)
+            print("\n==============================================\n")
+
+        try:
+            subprocess.check_call(gen_script, shell=True)
+        except subprocess.CalledProcessError as e:
+            files = os.listdir("./")
+            fstr = "\n-".join(files)
+            logging.error(f"Failed to Execute Script, return-code {e.returncode} \n"
+                          f"StdErr: {e.stderr}\n"
+                          f"StdOut: {e.stdout}\n"
+                          f" Current directory contents: .\n-{fstr}")
+            raise
+
+        final_outputs = []
+        for v in self._output_locs:
+            if issubclass(v.var_type, FlyteFile):
+                final_outputs.append(FlyteFile(outputs[v.var]))
+            if issubclass(v.var_type, FlyteDirectory):
+                final_outputs.append(FlyteDirectory(outputs[v.var]))
+        return tuple(final_outputs)
+
+    def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any:
+        return self._config_task_instance.post_execute(user_params, rval)
diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py
index d2ef84f2d2..a0112c1eea 100644
--- a/flytekit/types/directory/types.py
+++ b/flytekit/types/directory/types.py
@@ -172,14 +172,7 @@ def remote_source(self) -> str:
         return typing.cast(str, self._remote_source)
 
     def download(self) -> str:
-        if self._downloaded:
-            return self.path
-        if self._downloader is not noop:
-            self._downloader()
-            self._downloaded = True
-            return self.path
-        else:
-            raise ValueError(f"Attempting to trigger download on non-downloadable folder {self}")
+        return self.__fspath__()
 
     def __repr__(self):
         return self._path
diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py
index 9ebcc8e9a5..1fbcee049a 100644
--- a/flytekit/types/file/file.py
+++ b/flytekit/types/file/file.py
@@ -206,14 +206,7 @@ def remote_source(self) -> str:
         return typing.cast(str, self._remote_source)
 
     def download(self) -> str:
-        if self._downloaded:
-            return self.path
-        if self._downloader is not noop:
-            self._downloader()
-            self._downloaded = True
-            return self.path
-        else:
-            raise ValueError(f"Attempting to trigger download on non-downloadable file {self}")
+        return self.__fspath__()
 
     def __repr__(self):
         return self._path
diff --git a/tests/flytekit/unit/extras/tasks/__init__.py b/tests/flytekit/unit/extras/tasks/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
new file mode 100644
index 0000000000..18289bcee9
--- /dev/null
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -0,0 +1,76 @@
+import datetime
+
+import pytest
+
+from flytekit import kwtypes
+from flytekit.extras.tasks.shell import ShellTask, OutputLocation
+from flytekit.types.directory import FlyteDirectory
+from flytekit.types.file import CSVFile, FlyteFile
+
+
+def test_shell_task_no_io():
+    t = ShellTask(
+        name="test",
+        script="""
+        echo "Hello World!"
+        """,
+    )
+
+    t()
+
+
+def test_shell_task_fail():
+    t = ShellTask(
+        name="test",
+        script="""
+        non-existent blah
+        """,
+    )
+
+    with pytest.raises(Exception):
+        t()
+
+
+def test_input_substitution_primitive():
+    t = ShellTask(
+        name="test",
+        script="""
+        cat {{ .inputs.f }}
+        echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }}"
+        """,
+        inputs=kwtypes(f=str, y=int, j=datetime.datetime),
+    )
+
+    t(f="__init__.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    t(f="test_shell.py",  y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    with pytest.raises(Exception):
+        t(f="non_exist.py",  y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+
+
+def test_input_substitution_files():
+    t = ShellTask(
+        name="test",
+        script="""
+        cat {{ .inputs.f }}
+        echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }}"
+        """,
+        inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
+    )
+
+    t(f="testdata/test.csv",  y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+
+
+def test_input_output_substitution_files():
+    t = ShellTask(
+        name="test",
+        debug=True,
+        script="""
+        cat {{ .inputs.f }} >> {{ .outputs.y }}
+        echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"
+        """,
+        inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
+        output_locs=[OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
+                     OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")]
+    )
+
+    t(f="testdata/test.csv",  y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
\ No newline at end of file
diff --git a/tests/flytekit/unit/extras/tasks/testdata/test.csv b/tests/flytekit/unit/extras/tasks/testdata/test.csv
new file mode 100644
index 0000000000..e69de29bb2

From 6ae9446ddd804a96e14ab9eb9b9398d92989ca58 Mon Sep 17 00:00:00 2001
From: Ketan Umare <ketan.umare@gmail.com>
Date: Thu, 11 Nov 2021 17:16:34 -0800
Subject: [PATCH 2/6] Tests updated

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>
---
 flytekit/extras/tasks/shell.py                |  3 +-
 .../flytekit/unit/extras/tasks/test_shell.py  | 48 ++++++++++++++++---
 .../unit/extras/tasks/testdata/script.sh      |  4 ++
 3 files changed, 48 insertions(+), 7 deletions(-)
 create mode 100644 tests/flytekit/unit/extras/tasks/testdata/script.sh

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index d964173fe9..aae55112c7 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -90,7 +90,7 @@ def __init__(
             name: str,
             debug: bool = False,
             script: typing.Optional[str] = None,
-            script_file: typing.Optional[os.PathLike] = None,
+            script_file: typing.Optional[str] = None,
             task_config: T = None,
             inputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
             output_locs: typing.Optional[typing.List[OutputLocation]] = None,
@@ -114,6 +114,7 @@ def __init__(
         if script_file:
             if not os.path.exists(script_file):
                 raise ValueError(f"FileNotFound: the specified Script file at path {script_file} cannot be loaded")
+            script_file = os.path.abspath(script_file)
 
         # Each instance of NotebookTask instantiates an underlying task with a dummy function that will only be used
         # to run pre- and post- execute functions using the corresponding task plugin.
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 18289bcee9..ce4964cab5 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -1,4 +1,6 @@
 import datetime
+import os
+from subprocess import CalledProcessError
 
 import pytest
 
@@ -35,6 +37,7 @@ def test_input_substitution_primitive():
     t = ShellTask(
         name="test",
         script="""
+        set -ex
         cat {{ .inputs.f }}
         echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }}"
         """,
@@ -42,9 +45,9 @@ def test_input_substitution_primitive():
     )
 
     t(f="__init__.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
-    t(f="test_shell.py",  y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
-    with pytest.raises(Exception):
-        t(f="non_exist.py",  y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    t(f="test_shell.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    with pytest.raises(CalledProcessError):
+        t(f="non_exist.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
 
 
 def test_input_substitution_files():
@@ -57,15 +60,33 @@ def test_input_substitution_files():
         inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
     )
 
-    t(f="testdata/test.csv",  y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
 
 
 def test_input_output_substitution_files():
+    s = """
+        cat {{ .inputs.f }} >> {{ .outputs.y }}
+        echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"
+        """
+    t = ShellTask(
+        name="test",
+        debug=True,
+        script=s,
+        inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
+        output_locs=[OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
+                     OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")]
+    )
+
+    assert t.script == s
+    t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+
+
+def test_input_output_missing_var():
     t = ShellTask(
         name="test",
         debug=True,
         script="""
-        cat {{ .inputs.f }} >> {{ .outputs.y }}
+        cat {{ .inputs.f }} {{ .inputs.missing }} >> {{ .outputs.y }}
         echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"
         """,
         inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
@@ -73,4 +94,19 @@ def test_input_output_substitution_files():
                      OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")]
     )
 
-    t(f="testdata/test.csv",  y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
\ No newline at end of file
+    with pytest.raises(ValueError):
+        t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+
+
+def test_shell_script():
+    t = ShellTask(
+        name="test",
+        debug=True,
+        script_file="testdata/script.sh",
+        inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
+        output_locs=[OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
+                     OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")]
+    )
+
+    assert t.script_file == os.path.abspath("testdata/script.sh")
+    t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
diff --git a/tests/flytekit/unit/extras/tasks/testdata/script.sh b/tests/flytekit/unit/extras/tasks/testdata/script.sh
new file mode 100644
index 0000000000..ebe5aaf670
--- /dev/null
+++ b/tests/flytekit/unit/extras/tasks/testdata/script.sh
@@ -0,0 +1,4 @@
+set -ex
+
+cat {{ .inputs.f }} >> {{ .outputs.y }}
+echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"
\ No newline at end of file

From 6beda393ac6b11edcae84c1f612773c8d062a36f Mon Sep 17 00:00:00 2001
From: Ketan Umare <ketan.umare@gmail.com>
Date: Fri, 12 Nov 2021 14:05:40 -0800
Subject: [PATCH 3/6] lint fixed

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>
---
 flytekit/extras/tasks/shell.py                | 56 +++++++++++--------
 .../flytekit/unit/extras/tasks/test_shell.py  | 44 ++++++++++++---
 .../unit/extras/tasks/testdata/script.sh      |  6 +-
 3 files changed, 72 insertions(+), 34 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index aae55112c7..bbcd11e4bc 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -6,7 +6,7 @@
 import typing
 from dataclasses import dataclass
 
-from flytekit import PythonInstanceTask, ExecutionParameters
+from flytekit import ExecutionParameters, PythonInstanceTask
 from flytekit.core.interface import Interface
 from flytekit.core.task import TaskPlugins
 from flytekit.types.directory import FlyteDirectory
@@ -23,6 +23,7 @@ class OutputLocation:
                   vars and generates the path. Of the form ``"{{ .inputs.v }}.tmp.md"``.
                   This example for a given input v, at path `/tmp/abc.csv` will resolve to `/tmp/abc.csv.tmp.md`
     """
+
     var: str
     var_type: typing.Type
     location: typing.Union[os.PathLike, str]
@@ -46,8 +47,8 @@ def _stringify(v: typing.Any) -> str:
 
 def _interpolate(tmpl: str, regex: re.Pattern, validate_all_match: bool = True, **kwargs) -> str:
     """
-    Substitutes all templates that match the input regex - ``r"({{\s*.inputs.(\w+)\s*}})", re.IGNORECASE`` with the
-    given inputs and returns the substituted string. The result is non destructive towards the given string.
+    Substitutes all templates that match the supplied regex
+    with the given inputs and returns the substituted string. The result is non destructive towards the given string.
     """
     modified = tmpl
     matched = set()
@@ -79,22 +80,21 @@ def _dummy_task_func():
 
 
 class ShellTask(PythonInstanceTask[T]):
-    """
-    """
+    """ """
 
     _INPUT_REGEX = re.compile(r"({{\s*.inputs.(\w+)\s*}})", re.IGNORECASE)
     _OUTPUT_REGEX = re.compile(r"({{\s*.outputs.(\w+)\s*}})", re.IGNORECASE)
 
     def __init__(
-            self,
-            name: str,
-            debug: bool = False,
-            script: typing.Optional[str] = None,
-            script_file: typing.Optional[str] = None,
-            task_config: T = None,
-            inputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
-            output_locs: typing.Optional[typing.List[OutputLocation]] = None,
-            **kwargs,
+        self,
+        name: str,
+        debug: bool = False,
+        script: typing.Optional[str] = None,
+        script_file: typing.Optional[str] = None,
+        task_config: T = None,
+        inputs: typing.Optional[typing.Dict[str, typing.Type]] = None,
+        output_locs: typing.Optional[typing.List[OutputLocation]] = None,
+        **kwargs,
     ):
         """
         Args:
@@ -133,8 +133,11 @@ def __init__(
         self._output_locs = output_locs if output_locs else []
         outputs = self._validate_output_locs()
         super().__init__(
-            name, task_config, task_type=self._config_task_instance.task_type,
-            interface=Interface(inputs=inputs, outputs=outputs), **kwargs
+            name,
+            task_config,
+            task_type=self._config_task_instance.task_type,
+            interface=Interface(inputs=inputs, outputs=outputs),
+            **kwargs,
         )
 
     def _validate_output_locs(self) -> typing.Dict[str, typing.Type]:
@@ -147,8 +150,9 @@ def _validate_output_locs(self) -> typing.Dict[str, typing.Type]:
             if v.location is None:
                 raise ValueError(f"Output Location not provided for output var {v.var}")
             if not issubclass(v.var_type, FlyteFile) and not issubclass(v.var_type, FlyteDirectory):
-                raise ValueError(f"Currently only outputs of type FlyteFile/FlyteDirectory and their derived types"
-                                 f" are supported")
+                raise ValueError(
+                    "Currently only outputs of type FlyteFile/FlyteDirectory and their derived types are supported"
+                )
             outputs[v.var] = v.var_type
         return outputs
 
@@ -191,10 +195,12 @@ def execute(self, **kwargs) -> typing.Any:
         except subprocess.CalledProcessError as e:
             files = os.listdir("./")
             fstr = "\n-".join(files)
-            logging.error(f"Failed to Execute Script, return-code {e.returncode} \n"
-                          f"StdErr: {e.stderr}\n"
-                          f"StdOut: {e.stdout}\n"
-                          f" Current directory contents: .\n-{fstr}")
+            logging.error(
+                f"Failed to Execute Script, return-code {e.returncode} \n"
+                f"StdErr: {e.stderr}\n"
+                f"StdOut: {e.stdout}\n"
+                f" Current directory contents: .\n-{fstr}"
+            )
             raise
 
         final_outputs = []
@@ -203,7 +209,11 @@ def execute(self, **kwargs) -> typing.Any:
                 final_outputs.append(FlyteFile(outputs[v.var]))
             if issubclass(v.var_type, FlyteDirectory):
                 final_outputs.append(FlyteDirectory(outputs[v.var]))
-        return tuple(final_outputs)
+        if len(final_outputs) == 1:
+            return final_outputs[0]
+        if len(final_outputs) > 1:
+            return tuple(final_outputs)
+        return None
 
     def post_execute(self, user_params: ExecutionParameters, rval: typing.Any) -> typing.Any:
         return self._config_task_instance.post_execute(user_params, rval)
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index ce4964cab5..3c6f8f87f0 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -5,7 +5,7 @@
 import pytest
 
 from flytekit import kwtypes
-from flytekit.extras.tasks.shell import ShellTask, OutputLocation
+from flytekit.extras.tasks.shell import OutputLocation, ShellTask
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
@@ -60,7 +60,7 @@ def test_input_substitution_files():
         inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
     )
 
-    t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    assert t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0)) is None
 
 
 def test_input_output_substitution_files():
@@ -73,12 +73,34 @@ def test_input_output_substitution_files():
         debug=True,
         script=s,
         inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
-        output_locs=[OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
-                     OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")]
+        output_locs=[
+            OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
+            OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"),
+        ],
     )
 
     assert t.script == s
-    t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    x, y = t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    assert x is not None
+    assert y.path[-4:] == ".pyc"
+
+
+def test_input_single_output_substitution_files():
+    s = """
+        cat {{ .inputs.f }} >> {{ .outputs.y }}
+        echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }}"
+        """
+    t = ShellTask(
+        name="test",
+        debug=True,
+        script=s,
+        inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
+        output_locs=[OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")],
+    )
+
+    assert t.script == s
+    y = t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    assert y.path[-4:] == ".pyc"
 
 
 def test_input_output_missing_var():
@@ -90,8 +112,10 @@ def test_input_output_missing_var():
         echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"
         """,
         inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
-        output_locs=[OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
-                     OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")]
+        output_locs=[
+            OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
+            OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"),
+        ],
     )
 
     with pytest.raises(ValueError):
@@ -104,8 +128,10 @@ def test_shell_script():
         debug=True,
         script_file="testdata/script.sh",
         inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
-        output_locs=[OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
-                     OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")]
+        output_locs=[
+            OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
+            OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"),
+        ],
     )
 
     assert t.script_file == os.path.abspath("testdata/script.sh")
diff --git a/tests/flytekit/unit/extras/tasks/testdata/script.sh b/tests/flytekit/unit/extras/tasks/testdata/script.sh
index ebe5aaf670..22012ec3ae 100644
--- a/tests/flytekit/unit/extras/tasks/testdata/script.sh
+++ b/tests/flytekit/unit/extras/tasks/testdata/script.sh
@@ -1,4 +1,6 @@
+#!/bin/bash
+
 set -ex
 
-cat {{ .inputs.f }} >> {{ .outputs.y }}
-echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"
\ No newline at end of file
+cat "{{ .inputs.f }}" >> "{{ .outputs.y }}"
+echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"

From e383e7c174ab27776a8a9c7fbc0defa88bbfa6bf Mon Sep 17 00:00:00 2001
From: Ketan Umare <ketan.umare@gmail.com>
Date: Fri, 12 Nov 2021 14:23:59 -0800
Subject: [PATCH 4/6] tests updated

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>
---
 .../flytekit/unit/extras/tasks/test_shell.py  | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 3c6f8f87f0..002a037695 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -9,6 +9,10 @@
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
+testdata = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testdata")
+script_sh = os.path.join(testdata, "script.sh")
+test_csv = os.path.join(testdata, "test.csv")
+
 
 def test_shell_task_no_io():
     t = ShellTask(
@@ -60,7 +64,7 @@ def test_input_substitution_files():
         inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
     )
 
-    assert t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0)) is None
+    assert t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) is None
 
 
 def test_input_output_substitution_files():
@@ -80,7 +84,7 @@ def test_input_output_substitution_files():
     )
 
     assert t.script == s
-    x, y = t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    x, y = t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
     assert x is not None
     assert y.path[-4:] == ".pyc"
 
@@ -99,7 +103,7 @@ def test_input_single_output_substitution_files():
     )
 
     assert t.script == s
-    y = t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    y = t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
     assert y.path[-4:] == ".pyc"
 
 
@@ -119,14 +123,14 @@ def test_input_output_missing_var():
     )
 
     with pytest.raises(ValueError):
-        t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+        t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
 
 
 def test_shell_script():
     t = ShellTask(
-        name="test",
+        name="test2",
         debug=True,
-        script_file="testdata/script.sh",
+        script_file=script_sh,
         inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
         output_locs=[
             OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
@@ -134,5 +138,5 @@ def test_shell_script():
         ],
     )
 
-    assert t.script_file == os.path.abspath("testdata/script.sh")
-    t(f="testdata/test.csv", y="testdata", j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    assert t.script_file == script_sh
+    t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))

From 2f5edc07f8df6bc0c1a3c9f9750148796f16fe40 Mon Sep 17 00:00:00 2001
From: Ketan Umare <ketan.umare@gmail.com>
Date: Tue, 23 Nov 2021 21:47:30 -0800
Subject: [PATCH 5/6] updated

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>
---
 flytekit/extras/tasks/shell.py                |  2 +-
 .../flytekit/unit/extras/tasks/test_shell.py  | 28 ++++++++++++++++---
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index bbcd11e4bc..c6b0d20078 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -65,7 +65,7 @@ def _interpolate(tmpl: str, regex: re.Pattern, validate_all_match: bool = True,
     if validate_all_match:
         if len(matched) < len(kwargs.keys()):
             diff = set(kwargs.keys()).difference(matched)
-            raise ValueError(f"Extra Inputs have no matches in query template - missing {diff}")
+            raise ValueError(f"Extra Inputs have no matches in script template - missing {diff}")
     return modified
 
 
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index 002a037695..a8079baa1a 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -9,7 +9,8 @@
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import CSVFile, FlyteFile
 
-testdata = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testdata")
+test_file_path = os.path.dirname(os.path.realpath(__file__))
+testdata = os.path.join(test_file_path, "testdata")
 script_sh = os.path.join(testdata, "script.sh")
 test_csv = os.path.join(testdata, "test.csv")
 
@@ -48,8 +49,8 @@ def test_input_substitution_primitive():
         inputs=kwtypes(f=str, y=int, j=datetime.datetime),
     )
 
-    t(f="__init__.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
-    t(f="test_shell.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    t(f=os.path.join(test_file_path, "__init__.py"), y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+    t(f=os.path.join(test_file_path, "test_shell.py"), y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
     with pytest.raises(CalledProcessError):
         t(f="non_exist.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
 
@@ -107,7 +108,7 @@ def test_input_single_output_substitution_files():
     assert y.path[-4:] == ".pyc"
 
 
-def test_input_output_missing_var():
+def test_input_output_extra_var_in_template():
     t = ShellTask(
         name="test",
         debug=True,
@@ -126,6 +127,25 @@ def test_input_output_missing_var():
         t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
 
 
+def test_input_output_extra_input():
+    t = ShellTask(
+        name="test",
+        debug=True,
+        script="""
+        cat {{ .inputs.missing }} >> {{ .outputs.y }}
+        echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"
+        """,
+        inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
+        output_locs=[
+            OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
+            OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"),
+        ],
+    )
+
+    with pytest.raises(ValueError):
+        t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
+
+
 def test_shell_script():
     t = ShellTask(
         name="test2",

From d5b43902b9481bdcf0b5bfa137d016495a4c97e6 Mon Sep 17 00:00:00 2001
From: Ketan Umare <ketan.umare@gmail.com>
Date: Sun, 28 Nov 2021 12:20:42 -0800
Subject: [PATCH 6/6] better unit test

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>
---
 flytekit/extras/tasks/shell.py                |  9 +++++--
 .../flytekit/unit/extras/tasks/test_shell.py  | 25 +++++++++++++------
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py
index c6b0d20078..6e8dbcc21b 100644
--- a/flytekit/extras/tasks/shell.py
+++ b/flytekit/extras/tasks/shell.py
@@ -6,8 +6,9 @@
 import typing
 from dataclasses import dataclass
 
-from flytekit import ExecutionParameters, PythonInstanceTask
+from flytekit.core.context_manager import ExecutionParameters
 from flytekit.core.interface import Interface
+from flytekit.core.python_function_task import PythonInstanceTask
 from flytekit.core.task import TaskPlugins
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import FlyteFile
@@ -104,7 +105,7 @@ def __init__(
             script_file: A path to the file that contains the script (Only script or script_file) can be provided
             task_config: T Configuration for the task, can be either a Pod (or coming soon, BatchJob) config
             inputs: A Dictionary of input names to types
-            output_locs: A list of :ref:class:`OutputLocations`
+            output_locs: A list of :py:class:`OutputLocations`
             **kwargs: Other arguments that can be passed to :ref:class:`PythonInstanceTask`
         """
         if script and script_file:
@@ -116,6 +117,10 @@ def __init__(
                 raise ValueError(f"FileNotFound: the specified Script file at path {script_file} cannot be loaded")
             script_file = os.path.abspath(script_file)
 
+        if task_config is not None:
+            if str(type(task_config)) != "flytekitplugins.pod.task.Pod":
+                raise ValueError("TaskConfig can either be empty - indicating simple container task or a PodConfig.")
+
         # Each instance of NotebookTask instantiates an underlying task with a dummy function that will only be used
         # to run pre- and post- execute functions using the corresponding task plugin.
         # We rename the function name here to ensure the generated task has a unique name and avoid duplicate task name
diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py
index a8079baa1a..39f114a9c7 100644
--- a/tests/flytekit/unit/extras/tasks/test_shell.py
+++ b/tests/flytekit/unit/extras/tasks/test_shell.py
@@ -1,5 +1,6 @@
 import datetime
 import os
+import tempfile
 from subprocess import CalledProcessError
 
 import pytest
@@ -70,24 +71,32 @@ def test_input_substitution_files():
 
 def test_input_output_substitution_files():
     s = """
-        cat {{ .inputs.f }} >> {{ .outputs.y }}
-        echo "Hello World {{ .inputs.y }} on  {{ .inputs.j }} - output {{.outputs.x}}"
+        cat {{ .inputs.f }} > {{ .outputs.y }}
         """
     t = ShellTask(
         name="test",
         debug=True,
         script=s,
-        inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime),
+        inputs=kwtypes(f=CSVFile),
         output_locs=[
-            OutputLocation(var="x", var_type=FlyteDirectory, location="{{ .inputs.y }}"),
-            OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc"),
+            OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.mod"),
         ],
     )
 
     assert t.script == s
-    x, y = t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0))
-    assert x is not None
-    assert y.path[-4:] == ".pyc"
+
+    contents = "1,2,3,4\n"
+    with tempfile.TemporaryDirectory() as tmp:
+        csv = os.path.join(tmp, "abc.csv")
+        print(csv)
+        with open(csv, "w") as f:
+            f.write(contents)
+        y = t(f=csv)
+        assert y.path[-4:] == ".mod"
+        assert os.path.exists(y.path)
+        with open(y.path) as f:
+            s = f.read()
+        assert s == contents
 
 
 def test_input_single_output_substitution_files():