csv parser

Signed-off-by: Mandana Vaziri <[email protected]>

Author: vazirim

docs/tutorial.md

@@ -378,7 +378,7 @@ Bob lives at the following address:
 ## Parsing the output of a block
 
 As we saw in the previous section, it is possible to use the `parser: json` setting to parse the result of a block as a JSON.
-Other possible values for `parser` are `yaml`, `jsonl`, or `regex`.
+Other possible values for `parser` are `yaml`, `jsonl`, `regex`, or `csv`.
 
 The following example extracts using a regular expression parser the code between triple backtick generated by a model:
 

src/pdl/pdl-schema.json

@@ -3998,7 +3998,8 @@
           "enum": [
             "json",
             "jsonl",
-            "yaml"
+            "yaml",
+            "csv"
           ],
           "type": "string"
         },

src/pdl/pdl_ast.py

@@ -301,7 +301,7 @@ class RegexParser(Parser):
 
 
 ParserType = TypeAliasType(
-    "ParserType", Union[Literal["json", "jsonl", "yaml"], PdlParser, RegexParser]
+    "ParserType", Union[Literal["json", "jsonl", "yaml", "csv"], PdlParser, RegexParser]
 )
 """Different parsers."""
 OptionalParserType = TypeAliasType("OptionalParserType", Optional[ParserType])

src/pdl/pdl_interpreter.py

@@ -1,4 +1,5 @@
 # pylint: disable=import-outside-toplevel
+from io import StringIO
 import json
 import re
 import shlex
@@ -7,6 +8,7 @@
 import time
 import traceback
 import types
+import csv
 
 # TODO: temporarily disabling warnings to mute a pydantic warning from liteLLM
 import warnings
@@ -2647,6 +2649,18 @@ def parse_result(parser: ParserType, text: str) -> JSONReturnType:
                 raise PDLRuntimeParserError(
                     f"Attempted to parse ill-formed YAML: {repr(exc)}"
                 ) from exc
+        case "csv":
+            try:
+                result = []
+                reader = csv.reader(StringIO(text))
+                for row in reader:
+                    result.append(row)
+            except KeyboardInterrupt as exc:
+                raise exc from exc
+            except Exception as exc:
+                raise PDLRuntimeParserError(
+                    f"Attempted to parse ill-formed CSV: {repr(exc)}"
+                ) from exc
         case PdlParser():
             assert False, "TODO"
         case RegexParser(mode="search" | "match" | "fullmatch"):

tests/test_parser.py

@@ -172,3 +172,14 @@ def test_parser_case2():
     """
     result = exec_str(prog)
     assert result == ["1", "2", "3", "4"]
+
+def test_parser_csv():
+    csv_parser = """
+    text: |
+        1,Apple,Red
+        2,Orange,Orange
+        3,Banana,Yellow
+    parser: csv
+    """
+    result = exec_str(csv_parser)
+    assert result == [["1", "Apple", "Red"], ["2", "Orange", "Orange"], ["3", "Banana", "Yellow"]]