feat: add support for CodeBlocks with argv

starpit · starpit · commit d2d4b212ea60 · 2025-03-13T06:29:59.000-04:00
Signed-off-by: Nick Mitchell &lt;nickm@us.ibm.com&gt;
diff --git a/pdl-live-react/src/pdl_ast.d.ts b/pdl-live-react/src/pdl_ast.d.ts
@@ -1947,6 +1947,32 @@ export type Code =
   | ImportBlock
   | ErrorBlock
   | EmptyBlock
+  | (
+      | boolean
+      | number
+      | string
+      | FunctionBlock
+      | CallBlock
+      | LitellmModelBlock
+      | GraniteioModelBlock
+      | CodeBlock
+      | GetBlock
+      | DataBlock
+      | IfBlock
+      | MatchBlock
+      | RepeatBlock
+      | TextBlock
+      | LastOfBlock
+      | ArrayBlock
+      | ObjectBlock
+      | MessageBlock
+      | ReadBlock
+      | IncludeBlock
+      | ImportBlock
+      | ErrorBlock
+      | EmptyBlock
+      | null
+    )[]
   | null
 /**
  * Name of the variable used to store the result of the execution of the block.
diff --git a/pdl-live-react/src/pdl_ast_utils.ts b/pdl-live-react/src/pdl_ast_utils.ts
@@ -51,7 +51,7 @@ export function map_block_children(
       }
     })
     .with({ kind: "code" }, (block) => {
-      const code = f(block.code)
+      const code = Array.isArray(block.code) ? block.code.map(f) : f(block.code)
       return { ...block, code: code }
     })
     .with({ kind: "get" }, (block) => block)
@@ -155,7 +155,11 @@ export function iter_block_children(
       if (block.input) f(block.input)
     })
     .with({ kind: "code" }, (block) => {
-      f(block.code)
+      if (Array.isArray(block.code)) {
+        block.code.forEach(f)
+      } else {
+        f(block.code)
+      }
     })
     .with({ kind: "get" }, () => {})
     .with({ kind: "data" }, () => {})
diff --git a/src/pdl/pdl-schema.json b/src/pdl/pdl-schema.json
@@ -1380,6 +1380,88 @@
             {
               "$ref": "#/$defs/EmptyBlock"
             },
+            {
+              "items": {
+                "anyOf": [
+                  {
+                    "type": "boolean"
+                  },
+                  {
+                    "type": "integer"
+                  },
+                  {
+                    "type": "number"
+                  },
+                  {
+                    "type": "string"
+                  },
+                  {
+                    "$ref": "#/$defs/FunctionBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/CallBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/LitellmModelBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/GraniteioModelBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/CodeBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/GetBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/DataBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/IfBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/MatchBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/RepeatBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/TextBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/LastOfBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/ArrayBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/ObjectBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/MessageBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/ReadBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/IncludeBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/ImportBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/ErrorBlock"
+                  },
+                  {
+                    "$ref": "#/$defs/EmptyBlock"
+                  },
+                  {
+                    "type": "null"
+                  }
+                ]
+              },
+              "type": "array"
+            },
             {
               "type": "null"
             }
diff --git a/src/pdl/pdl_ast.py b/src/pdl/pdl_ast.py
@@ -15,7 +15,14 @@
     Union,
 )
 
-from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, RootModel
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    RootModel,
+    model_validator,
+)
 from pydantic.json_schema import SkipJsonSchema
 
 from .pdl_lazy import PdlDict, PdlLazy
@@ -462,10 +469,16 @@ class CodeBlock(LeafBlock):
     ]
     """Programming language of the code.
     """
-    code: "BlockType"
+    code: "BlockOrBlocksType"
     """Code to execute.
     """
 
+    @model_validator(mode="after")
+    def lang_is_python(self):
+        if isinstance(self.code, list) and self.lang != "command":
+            raise ValueError("CodeBlock code field is array for non-command block")
+        return self
+
 
 class GetBlock(LeafBlock):
     """
diff --git a/src/pdl/pdl_ast_utils.py b/src/pdl/pdl_ast_utils.py
@@ -48,7 +48,11 @@ def iter_block_children(f: Callable[[BlockType], None], block: BlockType) -> Non
             if block.pdl__trace is not None:
                 f(block.pdl__trace)
         case CodeBlock():
-            f(block.code)
+            if isinstance(block.code, list):
+                for b in block.code:
+                    f(b)
+            else:
+                f(block.code)
         case GetBlock():
             pass
         case DataBlock():
@@ -150,7 +154,10 @@ def map_block_children(f: MappedFunctions, block: BlockType) -> BlockType:
             if block.parameters is not None:
                 block.parameters = f.f_expr(block.parameters)
         case CodeBlock():
-            block.code = f.f_block(block.code)
+            if isinstance(block.code, list):
+                block.code = [f.f_block(b) for b in block.code]
+            else:
+                block.code = f.f_block(block.code)
         case GetBlock():
             pass
         case DataBlock():
diff --git a/src/pdl/pdl_dumper.py b/src/pdl/pdl_dumper.py
@@ -135,7 +135,10 @@ def block_to_dict(  # noqa: C901
                 d["modelResponse"] = block.modelResponse
         case CodeBlock():
             d["lang"] = block.lang
-            d["code"] = block_to_dict(block.code, json_compatible)
+            if isinstance(block.code, list):
+                d["code"] = [block_to_dict(b, json_compatible) for b in block.code]
+            else:
+                d["code"] = block_to_dict(block.code, json_compatible)
         case GetBlock():
             d["get"] = block.get
         case DataBlock():
diff --git a/src/pdl/pdl_interpreter.py b/src/pdl/pdl_interpreter.py
@@ -1431,14 +1431,24 @@ def process_call_code(
     state: InterpreterState, scope: ScopeType, block: CodeBlock, loc: PdlLocationType
 ) -> tuple[PdlLazy[Any], LazyMessages, ScopeType, CodeBlock]:
     background: LazyMessages
-    code_, _, _, block = process_block_of(
-        block,
-        "code",
-        state.with_yield_result(False).with_yield_background(False),
-        scope,
-        loc,
-    )
-    code_s = code_.result()
+    code_a = None
+    if isinstance(block.code, list):
+        code_s = ""
+        code_a, _, _, _ = process_block(
+            state.with_yield_result(False).with_yield_background(False),
+            scope,
+            ArrayBlock(array=block.code),
+            loc,
+        )
+    else:
+        code_, _, _, block = process_block_of(
+            block,
+            "code",
+            state.with_yield_result(False).with_yield_background(False),
+            scope,
+            loc,
+        )
+        code_s = code_.result()
     match block.lang:
         case "python":
             try:
@@ -1456,7 +1466,7 @@ def process_call_code(
                 ) from exc
         case "command":
             try:
-                result = call_command(code_s)
+                result = call_command(code_s, code_a)
                 background = PdlList(
                     [
                         PdlDict(  # type: ignore
@@ -1530,8 +1540,11 @@ def call_python(code: str, scope: ScopeType) -> PdlLazy[Any]:
     return PdlConst(result)
 
 
-def call_command(code: str) -> PdlLazy[str]:
-    args = shlex.split(code)
+def call_command(code: str, code_a: PdlLazy[list[str]] | None) -> PdlLazy[str]:
+    if code_a is not None and isinstance(code_a.result(), list):
+        args = code_a.result()
+    else:
+        args = shlex.split(code)
     p = subprocess.run(
         args, capture_output=True, text=True, check=False, shell=False
     )  # nosec B603
diff --git a/tests/test_code.py b/tests/test_code.py
@@ -61,6 +61,18 @@ def test_contribute_false():
     ]
 }
 
+command_data_args = {
+    "lastOf": [
+        {
+            "def": "world",
+            "lang": "command",
+            "code": ["echo", "-n", "World"],
+            "contribute": [],
+        },
+        "Hello ${ world }!",
+    ]
+}
+
 
 def test_command():
     result = exec_dict(command_data, output="all")
@@ -70,6 +82,14 @@ def test_command():
     assert scope["world"] == "World"
 
 
+def test_command_args():
+    result = exec_dict(command_data_args, output="all")
+    document = result["result"]
+    scope = result["scope"]
+    assert document == "Hello World!"
+    assert scope["world"] == "World"
+
+
 def test_jinja1():
     prog_str = """
 defs:
