change command to shell

patchwork/common/utils/utils.py

@@ -2,10 +2,14 @@
 
 import atexit
 import dataclasses
+import random
 import signal
+import string
 import tempfile
+from collections.abc import Mapping
 from pathlib import Path
 
+import chevron
 import tiktoken
 from chardet.universaldetector import UniversalDetector
 from git import Head, Repo
@@ -19,6 +23,20 @@
 _NEWLINES = {"\n", "\r\n", "\r"}
 
 
+def mustache_render(template: str, data: Mapping) -> str:
+    if len(data.keys()) < 1:
+        return template
+
+    chevron.render.__globals__["_html_escape"] = lambda x: x
+    return chevron.render(
+        template=template,
+        data=data,
+        partials_path=None,
+        partials_ext="".join(random.choices(string.ascii_uppercase + string.digits, k=32)),
+        partials_dict=dict(),
+    )
+
+
 def detect_newline(path: str | Path) -> str | None:
     with open(path, "r", newline="") as f:
         lines = f.read().splitlines(keepends=True)

patchwork/steps/CallSQL/CallSQL.py

@@ -2,14 +2,16 @@
 
 from sqlalchemy import URL, create_engine, exc, text
 
+from patchwork.common.utils.utils import mustache_render
 from patchwork.step import Step, StepStatus
 from patchwork.steps.CallSQL.typed import CallSQLInputs, CallSQLOutputs
 
 
 class CallSQL(Step, input_class=CallSQLInputs, output_class=CallSQLOutputs):
     def __init__(self, inputs: dict):
         super().__init__(inputs)
-        self.query = inputs["query"]
+        query_template_data = inputs.get("query_template_values", {})
+        self.query = mustache_render(inputs["query"], query_template_data)
         self.__build_engine(inputs)
 
     def __build_engine(self, inputs: dict):

patchwork/steps/CallSQL/typed.py

@@ -15,6 +15,7 @@ class CallSQLInputs(__RequiredCallSQLInputs, total=False):
     host: str
     port: int
     database: str
+    query_template_values: dict[str, Any]
 
 
 class CallSQLOutputs(TypedDict):

patchwork/steps/CallCommand/CallCommand.py → patchwork/steps/CallShell/CallShell.py

@@ -1,22 +1,20 @@
 from __future__ import annotations
 
 import shlex
-import shutil
 import subprocess
 from pathlib import Path
 
+from patchwork.common.utils.utils import mustache_render
 from patchwork.logger import logger
 from patchwork.step import Step, StepStatus
-from patchwork.steps.CallCommand.typed import CallCommandInputs, CallCommandOutputs
+from patchwork.steps.CallShell.typed import CallShellInputs, CallShellOutputs
 
 
-class CallCommand(Step, input_class=CallCommandInputs, output_class=CallCommandOutputs):
+class CallShell(Step, input_class=CallShellInputs, output_class=CallShellOutputs):
     def __init__(self, inputs: dict):
         super().__init__(inputs)
-        self.command = shutil.which(inputs["command"])
-        if self.command is None:
-            raise ValueError(f"Command `{inputs['command']}` not found in PATH")
-        self.command_args = shlex.split(inputs.get("command_args", ""))
+        script_template_values = inputs.get("script_template_values", {})
+        self.script = mustache_render(inputs["script"], script_template_values)
         self.working_dir = inputs.get("working_dir", Path.cwd())
         self.env = self.__parse_env_text(inputs.get("env", ""))
 
@@ -47,14 +45,12 @@ def __parse_env_text(env_text: str) -> dict[str, str]:
         return env
 
     def run(self) -> dict:
-        cmd = [self.command, *self.command_args]
-        p = subprocess.run(cmd, capture_output=True, text=True, cwd=self.working_dir, env=self.env)
+        p = subprocess.run(self.script, shell=True, capture_output=True, text=True, cwd=self.working_dir, env=self.env)
         try:
             p.check_returncode()
-            return dict(stdout_output=p.stdout)
         except subprocess.CalledProcessError as e:
             self.set_status(
                 StepStatus.FAILED,
-                f"`{self.command} {self.command_args}` failed with stdout:\n{p.stdout}\nstderr:\n{e.stderr}",
+                f"script failed with stdout:\n{p.stdout}\nstderr:\n{e.stderr}",
             )
-            return dict(stdout_output="")
+        return dict(stdout_output=p.stdout, stderr_output=p.stderr)

patchwork/steps/CallShell/typed.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing_extensions import Annotated, Any, TypedDict
+
+from patchwork.common.utils.step_typing import StepTypeConfig
+
+
+class __RequiredCallShellInputs(TypedDict):
+    script: str
+
+
+class CallShellInputs(__RequiredCallShellInputs, total=False):
+    working_dir: Annotated[str, StepTypeConfig(is_path=True)]
+    env: str
+    script_template_values: dict[str, Any]
+
+
+class CallShellOutputs(TypedDict):
+    stdout_output: str

patchwork/steps/PreparePrompt/PreparePrompt.py

@@ -1,13 +1,10 @@
 from __future__ import annotations
 
 import json
-import random
-import string
 from pathlib import Path
 
-import chevron
-
 from patchwork.common.constants import PROMPT_TEMPLATE_FILE_KEY
+from patchwork.common.utils.utils import mustache_render
 from patchwork.logger import logger
 from patchwork.step import Step, StepStatus
 
@@ -76,7 +73,6 @@ def run(self) -> dict:
             return dict(prompts=[])
 
         prompts = []
-        chevron.render.__globals__["_html_escape"] = lambda string: string
         for prompt_value in self.prompt_values:
             dict_value = prompt_value
             if not isinstance(dict_value, dict):
@@ -86,13 +82,7 @@ def run(self) -> dict:
             for prompt_part in self.prompt_template:
                 prompt_instance = {}
                 for key, value in prompt_part.items():
-                    new_value = chevron.render(
-                        template=value,
-                        data=dict_value,
-                        partials_path=None,
-                        partials_ext="".join(random.choices(string.ascii_uppercase + string.digits, k=32)),
-                        partials_dict=dict(),
-                    )
+                    new_value = mustache_render(value, dict_value)
                     prompt_instance[key] = new_value
                 prompt.append(prompt_instance)
             prompts.append(prompt)

patchwork/steps/__init__.py

@@ -3,6 +3,8 @@
 from patchwork.steps.CallAPI.CallAPI import CallAPI
 from patchwork.steps.CallCode2Prompt.CallCode2Prompt import CallCode2Prompt
 from patchwork.steps.CallLLM.CallLLM import CallLLM
+from patchwork.steps.CallShell.CallShell import CallShell
+from patchwork.steps.CallSQL.CallSQL import CallSQL
 from patchwork.steps.Combine.Combine import Combine
 from patchwork.steps.CommitChanges.CommitChanges import CommitChanges
 from patchwork.steps.CreateIssue.CreateIssue import CreateIssue
@@ -48,8 +50,6 @@
 from patchwork.steps.SimplifiedLLM.SimplifiedLLM import SimplifiedLLM
 from patchwork.steps.SimplifiedLLMOnce.SimplifiedLLMOnce import SimplifiedLLMOnce
 from patchwork.steps.SlackMessage.SlackMessage import SlackMessage
-from patchwork.steps.CallCommand.CallCommand import CallCommand
-from patchwork.steps.CallSQL.CallSQL import CallSQL
 
 # Compatibility Aliases
 JoinListPB = JoinList
@@ -63,7 +63,7 @@
     "AnalyzeImpact",
     "CallAPI",
     "CallCode2Prompt",
-    "CallCommand",
+    "CallShell",
     "CallSQL",
     "CallLLM",
     "Combine",