Ensure shebang lines are passed through as expected

examples/shebang-example.yaml

@@ -0,0 +1,22 @@
+# An example of steps that use bashisms, and declare that by way of a shebang line.
+# It would be better to put such a command in a script and call the script from the command field.
+- step:
+    name: Do a thing
+    image: python:3.12
+    command: |
+      #!/bin/bash
+      echo "The wibble is {parameter:wibble} today."
+      if [[ -z "$1" ]]; then echo "No arguments provided"; fi
+    parameters:
+      - name: wibble
+        type: integer
+- step:
+    name: Do the same thing in another way
+    image: python:3.12
+    command:
+      - "#!/bin/bash"
+      - echo "The wibble is {parameter:wibble} today."
+      - if [[ -z "$1" ]]; then echo "No arguments provided"; fi
+    parameters:
+      - name: wibble
+        type: integer

tests/conftest.py

@@ -23,3 +23,4 @@
 edge_merge_mode_config = config_fixture("edge-merge-mode.yaml")
 keep_order_config = config_fixture("keep-order.yaml")
 pipeline_with_retried_execution_config = config_fixture("pipeline-execution-retry-example.yaml")
+shebang_example_config = config_fixture("shebang-example.yaml")

tests/test_command.py

@@ -1,6 +1,6 @@
 from collections import defaultdict
 
-from valohai_yaml.commands import build_command
+from valohai_yaml.commands import build_command, join_command
 from valohai_yaml.objs.parameter_map import ParameterMap
 
 
@@ -130,3 +130,20 @@ def test_parameter_categories(example1_config):
         "Output": {"output-alias"},
         "Samples": {"unlabeled-samples", "labeled-samples"},
     }
+
+
+def test_shebang(shebang_example_config):
+    for step in shebang_example_config.steps.values():
+        command = step.build_command(parameter_values={"wibble": 42})
+        command = join_command(command, " && ")
+        assert command.startswith("#!/bin/bash\n")
+        assert "--wibble=42" in command
+
+
+def test_join_command():
+    assert join_command(["foo"], " && ") == "foo"
+    assert join_command(["foo", " ", "bar"], " && ") == "foo && bar"
+    assert join_command(["foo", " ", "bar"], "\n") == "foo\nbar"
+    assert join_command(["#!/bin/sh", "foo", " ", "bar"], " && ") == "#!/bin/sh\nfoo && bar"
+    assert join_command(["#!/bin/sh\nfoo", " ", "bar"], " && ") == "#!/bin/sh\nfoo && bar"
+    assert join_command(["#!/bin/bash\nfoo", "#!woop"], " && ") == "#!/bin/bash\nfoo && #!woop"

valohai_yaml/commands.py

@@ -58,7 +58,8 @@ def build_command(
 
     Even if passed a single `command`, this function will return a list
     of shell commands.  It is the caller's responsibility to concatenate them,
-    likely using the semicolon or double ampersands.
+    likely using the semicolon or double ampersands. The `join_command` function
+    can be used for this purpose.
 
     :param command: The command to interpolate params into.
     :param parameter_map: A ParameterMap object containing parameter knowledge.
@@ -94,3 +95,30 @@ def build_command(
                 )
         out_commands.append(command.strip())
     return out_commands
+
+
+def join_command(commands: List[str], joiner: str) -> str:
+    """
+    Join a list of commands into a single "script" that could be run using e.g. `sh -c`.
+
+    The commands (if not empty or purely whitespace)
+    are joined using the given joiner, but care is taken
+    to ensure that a shebang line, if any present at the start of the
+    first command, is preserved as a line of its own, as expected for
+    a script. (Note that `sh -c` wouldn't interpret it as a shebang line.)
+
+    :param commands: List of command pieces, e.g. from `build_command`.
+    :return: Single script string
+    """
+    shebang_line = None
+    bits = []
+    for i, cmd in enumerate(commands):
+        if i == 0 and cmd.startswith("#!"):
+            shebang_line, _, cmd = cmd.partition("\n")
+        if not cmd.strip():
+            continue
+        bits.append(cmd)
+    joined_bits = joiner.join(bits)
+    if shebang_line:
+        return f"{shebang_line}\n{joined_bits}"
+    return joined_bits