From 8b7c7a9261758529b9df79c8fdac12b09ac6de13 Mon Sep 17 00:00:00 2001 From: Zach Palchick Date: Wed, 22 Dec 2021 16:15:04 -0800 Subject: [PATCH 1/2] POC: Add support for f-string like sytax for shell task This commit is a proof of concept adding f-string like syntax for shell_tasks. This supports using nested types for script inputs, such as data classes. This change was motivated by the desire to combine shell_tasks that have multiple inputs with map_tasks which only support tasks with a single input. This commit is only a starting point, since it makes some changes to the shell_task API (adds a template_style field), and modifies some of the default behavior for ease of implementation (e.g. throwing an error when there are unused input arguments). Signed-off-by: Zach Palchick --- flytekit/extras/tasks/shell.py | 105 ++++++--- .../flytekit/unit/extras/tasks/test_shell.py | 207 ++++++++++++++---- 2 files changed, 239 insertions(+), 73 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 6e8dbcc21b..471f91b428 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -1,7 +1,10 @@ +import abc +import collections import datetime import logging import os import re +import string import subprocess import typing from dataclasses import dataclass @@ -46,30 +49,6 @@ def _stringify(v: typing.Any) -> str: return str(v) -def _interpolate(tmpl: str, regex: re.Pattern, validate_all_match: bool = True, **kwargs) -> str: - """ - 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() - 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 script template - missing {diff}") - return modified - - def _dummy_task_func(): """ A Fake function to satisfy the inner PythonTask requirements @@ -80,18 +59,81 @@ def _dummy_task_func(): T = typing.TypeVar("T") -class ShellTask(PythonInstanceTask[T]): - """ """ +class _Interpolaizer(abc.ABC): + @abc.abstractmethod + def interpolate(self, tmpl: str, inputs=None, outputs=None) -> str: + pass + + +class _DoubleCurlyBraceInterpolizer(_Interpolaizer): _INPUT_REGEX = re.compile(r"({{\s*.inputs.(\w+)\s*}})", re.IGNORECASE) _OUTPUT_REGEX = re.compile(r"({{\s*.outputs.(\w+)\s*}})", re.IGNORECASE) + def interpolate( + self, tmpl: str, inputs: typing.Dict[str, typing.Any] = None, outputs: typing.Dict[str, typing.Any] = None + ) -> str: + inputs = inputs or {} + outputs = outputs or {} + tmpl = self._interpolate(tmpl, self._INPUT_REGEX, **inputs) + tmpl = self._interpolate(tmpl, self._OUTPUT_REGEX, **outputs) + return tmpl + + @staticmethod + def _interpolate(tmpl: str, regex: re.Pattern, **kwargs) -> str: + """ + 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() + 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)) + return modified + + +class _PythonFStringInterpolizer(_Interpolaizer): + class _Formatter(string.Formatter): + def format_field(self, value, format_spec): + if isinstance(value, FlyteFile): + value.download() + return value.path + if isinstance(value, FlyteDirectory): + value.download() + return value.path + if isinstance(value, datetime.datetime): + return value.isoformat() + return super().format_field(value, format_spec) + + def interpolate(self, tmpl: str, inputs=None, outputs=None) -> str: + inputs = inputs or {} + outputs = outputs or {} + consolidated_args = collections.ChainMap(inputs, outputs) + try: + return self._Formatter().format(tmpl, **consolidated_args) + except KeyError as e: + raise ValueError(f"Variable {e} in Query not found in inputs {consolidated_args.keys()}") + + +class ShellTask(PythonInstanceTask[T]): + """ """ + + _interpolizers = {"default": _DoubleCurlyBraceInterpolizer, "python_f_string": _PythonFStringInterpolizer} + def __init__( self, name: str, debug: bool = False, script: typing.Optional[str] = None, script_file: typing.Optional[str] = None, + template_style: str = "default", task_config: T = None, inputs: typing.Optional[typing.Dict[str, typing.Type]] = None, output_locs: typing.Optional[typing.List[OutputLocation]] = None, @@ -136,6 +178,10 @@ def __init__( self._script_file = script_file self._debug = debug self._output_locs = output_locs if output_locs else [] + try: + self._interpolizer = self._interpolizers[template_style]() + except KeyError: + raise ValueError(f"'{template_style}' is not recognized as a valid template style") outputs = self._validate_output_locs() super().__init__( name, @@ -184,12 +230,9 @@ def execute(self, **kwargs) -> typing.Any: 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) + outputs[v.var] = self._interpolizer.interpolate(v.location, inputs=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) + gen_script = self._interpolizer.interpolate(self._script, inputs=kwargs, outputs=outputs) if self._debug: print("\n==============================================\n") print(gen_script) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index 39f114a9c7..d8f0e0c959 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -1,9 +1,11 @@ import datetime import os import tempfile +from dataclasses import dataclass from subprocess import CalledProcessError import pytest +from dataclasses_json import dataclass_json from flytekit import kwtypes from flytekit.extras.tasks.shell import OutputLocation, ShellTask @@ -39,14 +41,32 @@ def test_shell_task_fail(): t() -def test_input_substitution_primitive(): +@pytest.mark.parametrize( + "template_style,script", + [ + ( + "default", + """ + set -ex + cat {{ .inputs.f }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" + """, + ), + ( + "python_f_string", + """ + set -ex + cat {f} + echo "Hello World {y} on {j}" + """, + ), + ], +) +def test_input_substitution_primitive(template_style, script): t = ShellTask( name="test", - script=""" - set -ex - cat {{ .inputs.f }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """, + script=script, + template_style=template_style, inputs=kwtypes(f=str, y=int, j=datetime.datetime), ) @@ -56,34 +76,55 @@ def test_input_substitution_primitive(): t(f="non_exist.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) -def test_input_substitution_files(): +@pytest.mark.parametrize( + "template_style,script", + [ + ( + "default", + """ + cat {{ .inputs.f }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" + """, + ), + ( + "python_f_string", + """ + cat {f} + echo "Hello World {y} on {j}" + """, + ), + ], +) +def test_input_substitution_files(template_style, script): t = ShellTask( name="test", - script=""" - cat {{ .inputs.f }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """, + script=script, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), ) assert t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) is None -def test_input_output_substitution_files(): - s = """ - cat {{ .inputs.f }} > {{ .outputs.y }} - """ +@pytest.mark.parametrize( + "template_style,script,output_location", + [ + ("default", """cat {{ .inputs.f }} > {{ .outputs.y }}""", "{{ .inputs.f }}.mod"), + ("python_f_string", """cat {f} > {y}""", "{f}.mod"), + ], +) +def test_input_output_substitution_files(template_style, script, output_location): t = ShellTask( name="test", debug=True, - script=s, + script=script, + template_style=template_style, inputs=kwtypes(f=CSVFile), output_locs=[ - OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.mod"), + OutputLocation(var="y", var_type=FlyteFile, location=output_location), ], ) - assert t.script == s + assert t.script == script contents = "1,2,3,4\n" with tempfile.TemporaryDirectory() as tmp: @@ -99,62 +140,144 @@ def test_input_output_substitution_files(): assert s == contents -def test_input_single_output_substitution_files(): - s = """ - cat {{ .inputs.f }} >> {{ .outputs.y }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """ +@pytest.mark.parametrize( + "template_style,script,output_location", + [ + ( + "default", + """ + cat {{ .inputs.f }} >> {{ .outputs.z }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" + """, + "{{ .inputs.f }}.pyc", + ), + ( + "python_f_string", + """ + cat {f} >> {z} + echo "Hello World {y} on {j}" + """, + "{f}.pyc", + ), + ], +) +def test_input_single_output_substitution_files(template_style, script, output_location): t = ShellTask( name="test", debug=True, - script=s, + script=script, + template_style=template_style, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), - output_locs=[OutputLocation(var="y", var_type=FlyteFile, location="{{ .inputs.f }}.pyc")], + output_locs=[OutputLocation(var="z", var_type=FlyteFile, location=output_location)], ) - assert t.script == s + assert t.script == script y = t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) assert y.path[-4:] == ".pyc" -def test_input_output_extra_var_in_template(): +@pytest.mark.parametrize( + "template_style,script,output_location_x,output_location_z", + [ + ( + "default", + """ + cat {{ .inputs.f }} {{ .inputs.missing }} >> {{ .outputs.z }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" + """, + "{{ .inputs.y }}", + "{{ .inputs.f }}.pyc", + ), + ( + "python_f_string", + """ + cat {f} {missing} >> {z} + echo "Hello World {y} on {j} - output {x}" + """, + "{y}", + "{f}.pyc", + ), + ], +) +def test_input_output_extra_var_in_template(template_style, script, output_location_x, output_location_z): t = ShellTask( name="test", debug=True, - script=""" - cat {{ .inputs.f }} {{ .inputs.missing }} >> {{ .outputs.y }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" - """, + script=script, + template_style=template_style, 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"), + OutputLocation(var="x", var_type=FlyteDirectory, location=output_location_x), + OutputLocation(var="z", var_type=FlyteFile, location=output_location_z), ], ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="missing"): t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) -def test_input_output_extra_input(): +@pytest.mark.parametrize( + "template_style,script,output_location_x,output_location_z", + [ + ( + "default", + """ + cat {{ .inputs.missing }} >> {{ .outputs.y }} + echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" + """, + "{{ .inputs.y }}", + "{{ .inputs.f }}.pyc", + ), + ( + "python_f_string", + """ + cat {missing} >> {z} + echo "Hello World {y} on {j} - output {x}" + """, + "{y}", + "{f}.pyc", + ), + ], +) +def test_input_output_extra_input(template_style, script, output_location_x, output_location_z): t = ShellTask( name="test", debug=True, - script=""" - cat {{ .inputs.missing }} >> {{ .outputs.y }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" - """, + script=script, + template_style=template_style, 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"), + OutputLocation(var="x", var_type=FlyteDirectory, location=output_location_x), + OutputLocation(var="z", var_type=FlyteFile, location=output_location_z), ], ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="missing"): t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) +def test_can_use_complex_types_for_inputs_to_f_string_template(): + @dataclass_json + @dataclass + class InputArgs: + in_file: CSVFile + + t = ShellTask( + name="test", + debug=True, + script="""cat {input_args.in_file} >> {input_args.in_file}.tmp""", + template_style="python_f_string", + inputs=kwtypes(input_args=InputArgs), + output_locs=[ + OutputLocation(var="x", var_type=FlyteFile, location="{input_args.in_file}.tmp"), + ], + ) + + input_args = InputArgs(FlyteFile(path=test_csv)) + x = t(input_args=input_args) + assert x.path[-4:] == ".tmp" + + def test_shell_script(): t = ShellTask( name="test2", From 92acb37a1a202852743d5640972e7783a2022c28 Mon Sep 17 00:00:00 2001 From: Zach Palchick Date: Tue, 4 Jan 2022 10:32:32 -0800 Subject: [PATCH 2/2] Drop support for old/regex style for doing string interpolation Signed-off-by: Zach Palchick --- flytekit/extras/tasks/shell.py | 87 +++------- .../flytekit/unit/extras/tasks/test_shell.py | 153 ++++-------------- .../unit/extras/tasks/testdata/script.sh | 4 +- 3 files changed, 55 insertions(+), 189 deletions(-) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 471f91b428..63bae594ed 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -1,9 +1,7 @@ -import abc import collections import datetime import logging import os -import re import string import subprocess import typing @@ -33,22 +31,6 @@ class OutputLocation: 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 _dummy_task_func(): """ A Fake function to satisfy the inner PythonTask requirements @@ -59,49 +41,16 @@ def _dummy_task_func(): T = typing.TypeVar("T") -class _Interpolaizer(abc.ABC): - @abc.abstractmethod - def interpolate(self, tmpl: str, inputs=None, outputs=None) -> str: - pass - - -class _DoubleCurlyBraceInterpolizer(_Interpolaizer): - - _INPUT_REGEX = re.compile(r"({{\s*.inputs.(\w+)\s*}})", re.IGNORECASE) - _OUTPUT_REGEX = re.compile(r"({{\s*.outputs.(\w+)\s*}})", re.IGNORECASE) - - def interpolate( - self, tmpl: str, inputs: typing.Dict[str, typing.Any] = None, outputs: typing.Dict[str, typing.Any] = None - ) -> str: - inputs = inputs or {} - outputs = outputs or {} - tmpl = self._interpolate(tmpl, self._INPUT_REGEX, **inputs) - tmpl = self._interpolate(tmpl, self._OUTPUT_REGEX, **outputs) - return tmpl +class _PythonFStringInterpolizer: + """A class for interpolating scripts that use python string.format syntax""" - @staticmethod - def _interpolate(tmpl: str, regex: re.Pattern, **kwargs) -> str: - """ - 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() - 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)) - return modified - - -class _PythonFStringInterpolizer(_Interpolaizer): class _Formatter(string.Formatter): def format_field(self, value, format_spec): + """ + 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(value, FlyteFile): value.download() return value.path @@ -112,9 +61,21 @@ def format_field(self, value, format_spec): return value.isoformat() return super().format_field(value, format_spec) - def interpolate(self, tmpl: str, inputs=None, outputs=None) -> str: + def interpolate( + self, + tmpl: str, + inputs: typing.Optional[typing.Dict[str, str]] = None, + outputs: typing.Optional[typing.Dict[str, str]] = None, + ) -> str: + """ + Interpolate python formatted string templates with variables from the input and output + argument dicts. The result is non destructive towards the given template string. + """ inputs = inputs or {} outputs = outputs or {} + reused_vars = inputs.keys() & outputs.keys() + if reused_vars: + raise ValueError(f"Variables {reused_vars} in Query cannot be shared between inputs and outputs.") consolidated_args = collections.ChainMap(inputs, outputs) try: return self._Formatter().format(tmpl, **consolidated_args) @@ -125,15 +86,12 @@ def interpolate(self, tmpl: str, inputs=None, outputs=None) -> str: class ShellTask(PythonInstanceTask[T]): """ """ - _interpolizers = {"default": _DoubleCurlyBraceInterpolizer, "python_f_string": _PythonFStringInterpolizer} - def __init__( self, name: str, debug: bool = False, script: typing.Optional[str] = None, script_file: typing.Optional[str] = None, - template_style: str = "default", task_config: T = None, inputs: typing.Optional[typing.Dict[str, typing.Type]] = None, output_locs: typing.Optional[typing.List[OutputLocation]] = None, @@ -178,10 +136,7 @@ def __init__( self._script_file = script_file self._debug = debug self._output_locs = output_locs if output_locs else [] - try: - self._interpolizer = self._interpolizers[template_style]() - except KeyError: - raise ValueError(f"'{template_style}' is not recognized as a valid template style") + self._interpolizer = _PythonFStringInterpolizer() outputs = self._validate_output_locs() super().__init__( name, diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index d8f0e0c959..2b76f7ad7f 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -41,32 +41,14 @@ def test_shell_task_fail(): t() -@pytest.mark.parametrize( - "template_style,script", - [ - ( - "default", - """ - set -ex - cat {{ .inputs.f }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """, - ), - ( - "python_f_string", - """ +def test_input_substitution_primitive(): + t = ShellTask( + name="test", + script=""" set -ex cat {f} echo "Hello World {y} on {j}" """, - ), - ], -) -def test_input_substitution_primitive(template_style, script): - t = ShellTask( - name="test", - script=script, - template_style=template_style, inputs=kwtypes(f=str, y=int, j=datetime.datetime), ) @@ -76,51 +58,28 @@ def test_input_substitution_primitive(template_style, script): t(f="non_exist.py", y=5, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) -@pytest.mark.parametrize( - "template_style,script", - [ - ( - "default", - """ - cat {{ .inputs.f }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """, - ), - ( - "python_f_string", - """ +def test_input_substitution_files(): + t = ShellTask( + name="test", + script=""" cat {f} echo "Hello World {y} on {j}" """, - ), - ], -) -def test_input_substitution_files(template_style, script): - t = ShellTask( - name="test", - script=script, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), ) assert t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) is None -@pytest.mark.parametrize( - "template_style,script,output_location", - [ - ("default", """cat {{ .inputs.f }} > {{ .outputs.y }}""", "{{ .inputs.f }}.mod"), - ("python_f_string", """cat {f} > {y}""", "{f}.mod"), - ], -) -def test_input_output_substitution_files(template_style, script, output_location): +def test_input_output_substitution_files(): + script = "cat {f} > {y}" t = ShellTask( name="test", debug=True, script=script, - template_style=template_style, inputs=kwtypes(f=CSVFile), output_locs=[ - OutputLocation(var="y", var_type=FlyteFile, location=output_location), + OutputLocation(var="y", var_type=FlyteFile, location="{f}.mod"), ], ) @@ -140,35 +99,17 @@ def test_input_output_substitution_files(template_style, script, output_location assert s == contents -@pytest.mark.parametrize( - "template_style,script,output_location", - [ - ( - "default", - """ - cat {{ .inputs.f }} >> {{ .outputs.z }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }}" - """, - "{{ .inputs.f }}.pyc", - ), - ( - "python_f_string", - """ +def test_input_single_output_substitution_files(): + script = """ cat {f} >> {z} echo "Hello World {y} on {j}" - """, - "{f}.pyc", - ), - ], -) -def test_input_single_output_substitution_files(template_style, script, output_location): + """ t = ShellTask( name="test", debug=True, script=script, - template_style=template_style, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), - output_locs=[OutputLocation(var="z", var_type=FlyteFile, location=output_location)], + output_locs=[OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc")], ) assert t.script == script @@ -177,38 +118,31 @@ def test_input_single_output_substitution_files(template_style, script, output_l @pytest.mark.parametrize( - "template_style,script,output_location_x,output_location_z", + "script", [ ( - "default", """ - cat {{ .inputs.f }} {{ .inputs.missing }} >> {{ .outputs.z }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" - """, - "{{ .inputs.y }}", - "{{ .inputs.f }}.pyc", + cat {missing} >> {z} + echo "Hello World {y} on {j} - output {x}" + """ ), ( - "python_f_string", """ cat {f} {missing} >> {z} echo "Hello World {y} on {j} - output {x}" - """, - "{y}", - "{f}.pyc", + """ ), ], ) -def test_input_output_extra_var_in_template(template_style, script, output_location_x, output_location_z): +def test_input_output_extra_and_missing_variables(script): t = ShellTask( name="test", debug=True, script=script, - template_style=template_style, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), output_locs=[ - OutputLocation(var="x", var_type=FlyteDirectory, location=output_location_x), - OutputLocation(var="z", var_type=FlyteFile, location=output_location_z), + OutputLocation(var="x", var_type=FlyteDirectory, location="{y}"), + OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc"), ], ) @@ -216,43 +150,21 @@ def test_input_output_extra_var_in_template(template_style, script, output_locat t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) -@pytest.mark.parametrize( - "template_style,script,output_location_x,output_location_z", - [ - ( - "default", - """ - cat {{ .inputs.missing }} >> {{ .outputs.y }} - echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" - """, - "{{ .inputs.y }}", - "{{ .inputs.f }}.pyc", - ), - ( - "python_f_string", - """ - cat {missing} >> {z} - echo "Hello World {y} on {j} - output {x}" - """, - "{y}", - "{f}.pyc", - ), - ], -) -def test_input_output_extra_input(template_style, script, output_location_x, output_location_z): +def test_cannot_reuse_variables_for_both_inputs_and_outputs(): t = ShellTask( name="test", debug=True, - script=script, - template_style=template_style, + script=""" + cat {f} >> {y} + echo "Hello World {y} on {j}" + """, inputs=kwtypes(f=CSVFile, y=FlyteDirectory, j=datetime.datetime), output_locs=[ - OutputLocation(var="x", var_type=FlyteDirectory, location=output_location_x), - OutputLocation(var="z", var_type=FlyteFile, location=output_location_z), + OutputLocation(var="y", var_type=FlyteFile, location="{f}.pyc"), ], ) - with pytest.raises(ValueError, match="missing"): + with pytest.raises(ValueError, match="Variables {'y'} in Query"): t(f=test_csv, y=testdata, j=datetime.datetime(2021, 11, 10, 12, 15, 0)) @@ -266,7 +178,6 @@ class InputArgs: name="test", debug=True, script="""cat {input_args.in_file} >> {input_args.in_file}.tmp""", - template_style="python_f_string", inputs=kwtypes(input_args=InputArgs), output_locs=[ OutputLocation(var="x", var_type=FlyteFile, location="{input_args.in_file}.tmp"), @@ -285,8 +196,8 @@ def test_shell_script(): script_file=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"), + OutputLocation(var="x", var_type=FlyteDirectory, location="{y}"), + OutputLocation(var="z", var_type=FlyteFile, location="{f}.pyc"), ], ) diff --git a/tests/flytekit/unit/extras/tasks/testdata/script.sh b/tests/flytekit/unit/extras/tasks/testdata/script.sh index 22012ec3ae..1deb4c474a 100644 --- a/tests/flytekit/unit/extras/tasks/testdata/script.sh +++ b/tests/flytekit/unit/extras/tasks/testdata/script.sh @@ -2,5 +2,5 @@ set -ex -cat "{{ .inputs.f }}" >> "{{ .outputs.y }}" -echo "Hello World {{ .inputs.y }} on {{ .inputs.j }} - output {{.outputs.x}}" +cat "{f}" >> "{z}" +echo "Hello World {y} on {j} - output {x}"