Fix inline script parsing issues (#574)

Fixes #572 The fix makes the function parser a bit more robust by round-tripping through ast parse/unparse. This ensures we get consistently formatted code each time and removes any extraneous comments which were previously causing parsing issues. Signed-off-by: Sambhav Kothari <skothari44@bloomberg.net>
argoproj-labs · Apr 15, 2023 · 7f6bec0 · 7f6bec0
1 parent 3fa4cc8
commit 7f6bec0
Show file tree

Hide file tree

Showing 49 changed files with 1,195 additions and 232 deletions.
diff --git a/docs/examples/workflows/coinflip.md b/docs/examples/workflows/coinflip.md
@@ -72,12 +72,9 @@
 
             import random
 
+            result = ''heads'' if random.randint(0, 1) == 0 else ''tails''
 
-            result = "heads" if random.randint(0, 1) == 0 else "tails"
-
-            print(result)
-
-            '
+            print(result)'
       - name: heads
         script:
           command:
@@ -89,9 +86,7 @@
 
             sys.path.append(os.getcwd())
 
-            print("it was heads")
-
-            '
+            print(''it was heads'')'
       - name: tails
         script:
           command:
@@ -103,8 +98,6 @@
 
             sys.path.append(os.getcwd())
 
-            print("it was tails")
-
-            '
+            print(''it was tails'')'
     ```
 
diff --git a/docs/examples/workflows/complex_deps.md b/docs/examples/workflows/complex_deps.md
@@ -78,6 +78,6 @@
           image: python:3.8
           source: "import os\nimport sys\nsys.path.append(os.getcwd())\nimport json\n\
             try: p = json.loads(r'''{{inputs.parameters.p}}''')\nexcept: p = r'''{{inputs.parameters.p}}'''\n\
-            \nif p < 0.5:\n    raise Exception(p)\nprint(42)\n"
+            \nif p < 0.5:\n    raise Exception(p)\nprint(42)"
     ```
 
diff --git a/docs/examples/workflows/dag-with-script-output-param-passing.md b/docs/examples/workflows/dag-with-script-output-param-passing.md
@@ -68,8 +68,8 @@
           command:
           - python
           image: python:3.8
-          source: "import os\nimport sys\nsys.path.append(os.getcwd())\nwith open(\"/test\"\
-            , \"w\") as f_out:\n    f_out.write(\"test\")\n"
+          source: "import os\nimport sys\nsys.path.append(os.getcwd())\nwith open('/test',\
+            \ 'w') as f_out:\n    f_out.write('test')"
       - inputs:
           parameters:
           - name: a
@@ -91,8 +91,6 @@
             except: a = r''''''{{inputs.parameters.a}}''''''
 
 
-            print(a)
-
-            '
+            print(a)'
     ```
 
diff --git a/docs/examples/workflows/dag_conditional_parameters.md b/docs/examples/workflows/dag_conditional_parameters.md
@@ -79,25 +79,18 @@
           image: python:alpine3.6
           source: 'import random
 
-
-            print("heads" if random.randint(0, 1) == 0 else "tails")
-
-            '
+            print(''heads'' if random.randint(0, 1) == 0 else ''tails'')'
       - name: heads
         script:
           command:
           - python
           image: python:alpine3.6
-          source: 'print("heads")
-
-            '
+          source: print('heads')
       - name: tails
         script:
           command:
           - python
           image: python:alpine3.6
-          source: 'print("tails")
-
-            '
+          source: print('tails')
     ```
 
diff --git a/docs/examples/workflows/dag_diamond_with_callable_decorators.md b/docs/examples/workflows/dag_diamond_with_callable_decorators.md
@@ -84,8 +84,6 @@
             except: message = r''''''{{inputs.parameters.message}}''''''
 
 
-            print(message)
-
-            '
+            print(message)'
     ```
 
diff --git a/docs/examples/workflows/dag_diamond_with_callable_script.md b/docs/examples/workflows/dag_diamond_with_callable_script.md
@@ -70,9 +70,7 @@
             except: message = r''''''{{inputs.parameters.message}}''''''
 
 
-            print(message)
-
-            '
+            print(message)'
       - dag:
           tasks:
           - arguments:

diff --git a/docs/examples/workflows/dag_with_script_param_passing.md b/docs/examples/workflows/dag_with_script_param_passing.md
@@ -61,9 +61,7 @@
 
             sys.path.append(os.getcwd())
 
-            print(42)
-
-            '
+            print(42)'
       - inputs:
           parameters:
           - name: a
@@ -85,8 +83,6 @@
             except: a = r''''''{{inputs.parameters.a}}''''''
 
 
-            print(a)
-
-            '
+            print(a)'
     ```
 
diff --git a/docs/examples/workflows/dynamic_volumes.md b/docs/examples/workflows/dynamic_volumes.md
@@ -51,10 +51,7 @@
 
             import subprocess
 
-
-            print(subprocess.run("cd && /mnt && df -h", shell=True, capture_output=True).stdout.decode())
-
-            '
+            print(subprocess.run(''cd && /mnt && df -h'', shell=True, capture_output=True).stdout.decode())'
           volumeMounts:
           - mountPath: /mnt/vol
             name: v

diff --git a/docs/examples/workflows/global_config.md b/docs/examples/workflows/global_config.md
@@ -57,8 +57,6 @@
 
             sys.path.append(os.getcwd())
 
-            print("hello")
-
-            '
+            print(''hello'')'
     ```
 
diff --git a/docs/examples/workflows/multi_env.md b/docs/examples/workflows/multi_env.md
@@ -61,15 +61,10 @@
 
             import os
 
+            assert os.environ[''a''] == ''1'', os.environ[''a'']
 
-            # note that env params come in as strings
+            assert os.environ[''b''] == ''2'', os.environ[''b'']
 
-            assert os.environ["a"] == "1", os.environ["a"]
-
-            assert os.environ["b"] == "2", os.environ["b"]
-
-            assert os.environ["c"] == "3", os.environ["c"]
-
-            '
+            assert os.environ[''c''] == ''3'', os.environ[''c'']'
     ```
 
diff --git a/docs/examples/workflows/script_auto_infer.md b/docs/examples/workflows/script_auto_infer.md
@@ -68,8 +68,8 @@
           - python
           image: python:3.8
           source: "import os\nimport sys\nsys.path.append(os.getcwd())\nimport pickle\n\
-            \nresult = \"foo testing\"\nwith open(\"/tmp/result\", \"wb\") as f:\n   \
-            \ pickle.dump(result, f)\n"
+            result = 'foo testing'\nwith open('/tmp/result', 'wb') as f:\n    pickle.dump(result,\
+            \ f)"
       - inputs:
           artifacts:
           - name: i
@@ -80,7 +80,6 @@
           - python
           image: python:3.8
           source: "import os\nimport sys\nsys.path.append(os.getcwd())\nimport json\n\n\
-            import pickle\n\nwith open(\"/tmp/i\", \"rb\") as f:\n    i = pickle.load(f)\n\
-            print(i)\n"
+            import pickle\nwith open('/tmp/i', 'rb') as f:\n    i = pickle.load(f)\nprint(i)"
     ```
 
diff --git a/docs/examples/workflows/script_variations.md b/docs/examples/workflows/script_variations.md
@@ -0,0 +1,82 @@
+# Script Variations
+
+
+
+
+
+
+=== "Hera"
+
+    ```python linenums="1"
+    from hera.workflows import Workflow, script
+
+
+    @script()
+    def hello_world():  # pragma: no cover
+        print("Hello World!")
+
+
+    @script()
+    def multiline_function(
+        test: str,
+        another_test: str,
+    ) -> str:  # pragma: no cover
+        print("Hello World!")
+
+
+    with Workflow(generate_name="fv-test-", entrypoint="d") as w:
+        hello_world()
+        multiline_function()
+    ```
+
+=== "YAML"
+
+    ```yaml linenums="1"
+    apiVersion: argoproj.io/v1alpha1
+    kind: Workflow
+    metadata:
+      generateName: fv-test-
+    spec:
+      entrypoint: d
+      templates:
+      - name: hello-world
+        script:
+          command:
+          - python
+          image: python:3.8
+          source: 'import os
+
+            import sys
+
+            sys.path.append(os.getcwd())
+
+            print(''Hello World!'')'
+      - inputs:
+          parameters:
+          - name: test
+          - name: another_test
+        name: multiline-function
+        script:
+          command:
+          - python
+          image: python:3.8
+          source: 'import os
+
+            import sys
+
+            sys.path.append(os.getcwd())
+
+            import json
+
+            try: another_test = json.loads(r''''''{{inputs.parameters.another_test}}'''''')
+
+            except: another_test = r''''''{{inputs.parameters.another_test}}''''''
+
+            try: test = json.loads(r''''''{{inputs.parameters.test}}'''''')
+
+            except: test = r''''''{{inputs.parameters.test}}''''''
+
+
+            print(''Hello World!'')'
+    ```
+
diff --git a/docs/examples/workflows/script_with_default_params.md b/docs/examples/workflows/script_with_default_params.md
@@ -102,8 +102,6 @@
             except: c = r''''''{{inputs.parameters.c}}''''''
 
 
-            print(a, b, c)
-
-            '
+            print(a, b, c)'
     ```
 
diff --git a/docs/examples/workflows/script_with_resources.md b/docs/examples/workflows/script_with_resources.md
@@ -44,8 +44,6 @@
 
             sys.path.append(os.getcwd())
 
-            print("ok")
-
-            '
+            print(''ok'')'
     ```
 
diff --git a/docs/examples/workflows/upstream/coinflip.md b/docs/examples/workflows/upstream/coinflip.md
@@ -95,11 +95,8 @@
           image: python:alpine3.6
           source: 'import random
 
+            result = ''heads'' if random.randint(0, 1) == 0 else ''tails''
 
-            result = "heads" if random.randint(0, 1) == 0 else "tails"
-
-            print(result)
-
-            '
+            print(result)'
     ```
 
diff --git a/docs/examples/workflows/upstream/colored_logs.md b/docs/examples/workflows/upstream/colored_logs.md
@@ -53,12 +53,11 @@
           - name: PYTHONUNBUFFERED
             value: '1'
           image: python:3.7
-          source: "import time  # noqa: I001\nimport random\n\nmessages = [\n    \"No\
-            \ Color\",\n    \"\\x1b[30m%s\\x1b[0m\" % \"FG Black\",\n    \"\\x1b[32m%s\\\
-            x1b[0m\" % \"FG Green\",\n    \"\\x1b[34m%s\\x1b[0m\" % \"FG Blue\",\n   \
-            \ \"\\x1b[36m%s\\x1b[0m\" % \"FG Cyan\",\n    \"\\x1b[41m%s\\x1b[0m\" % \"\
-            BG Red\",\n    \"\\x1b[43m%s\\x1b[0m\" % \"BG Yellow\",\n    \"\\x1b[45m%s\\\
-            x1b[0m\" % \"BG Magenta\",\n]\nfor i in range(1, 100):\n    print(random.choice(messages))\n\
-            \    time.sleep(1)\n"
+          source: "import time\nimport random\nmessages = ['No Color', '\\x1b[30m%s\\\
+            x1b[0m' % 'FG Black', '\\x1b[32m%s\\x1b[0m' % 'FG Green', '\\x1b[34m%s\\x1b[0m'\
+            \ % 'FG Blue', '\\x1b[36m%s\\x1b[0m' % 'FG Cyan', '\\x1b[41m%s\\x1b[0m' %\
+            \ 'BG Red', '\\x1b[43m%s\\x1b[0m' % 'BG Yellow', '\\x1b[45m%s\\x1b[0m' % 'BG\
+            \ Magenta']\nfor i in range(1, 100):\n    print(random.choice(messages))\n\
+            \    time.sleep(1)"
     ```
 
diff --git a/docs/examples/workflows/upstream/conditional_artifacts.md b/docs/examples/workflows/upstream/conditional_artifacts.md
@@ -106,10 +106,7 @@
           image: python:alpine3.6
           source: 'import random
 
-
-            print("heads" if random.randint(0, 1) == 0 else "tails")
-
-            '
+            print(''heads'' if random.randint(0, 1) == 0 else ''tails'')'
       - name: heads
         outputs:
           artifacts:
@@ -119,8 +116,7 @@
           command:
           - python
           image: python:alpine3.6
-          source: "with open(\"result.txt\", \"w\") as f:\n    f.write(\"it was heads\"\
-            )\n"
+          source: "with open('result.txt', 'w') as f:\n    f.write('it was heads')"
       - name: tails
         outputs:
           artifacts:
@@ -130,8 +126,7 @@
           command:
           - python
           image: python:alpine3.6
-          source: "with open(\"result.txt\", \"w\") as f:\n    f.write(\"it was tails\"\
-            )\n"
+          source: "with open('result.txt', 'w') as f:\n    f.write('it was tails')"
       - name: main
         outputs:
           artifacts:

diff --git a/.../examples/workflows/upstream/container_set_template__outputs_result_workflow.md b/.../examples/workflows/upstream/container_set_template__outputs_result_workflow.md
@@ -80,9 +80,7 @@
           command:
           - python
           image: python:alpine3.6
-          source: 'assert "{{inputs.parameters.x}}" == "hi"
-
-            '
+          source: assert '{{inputs.parameters.x}}' == 'hi'
       - dag:
           tasks:
           - name: a