(DOP-2651) Implement input, output, and io-code-block directives (#375)

* (DOP-2651) Implement input, output, and io-code-block directives

* add parser tests and new source option

* add logic and tests for handling edge cases

* fix formatting and make functions for repeated code

* fix breaking tests and reformat code

* fix structuring and simplifying logic

* add language as an option for input/output, remove it as an argument for io-code-block

* fix fstring issue

* comment out examples

* add arg to rstspec.toml for error checking purposes

* fix docstrings

Co-authored-by: Grace Chong <[email protected]>

Author: gjchong25

snooty/diagnostics.py

@@ -180,6 +180,18 @@ def __init__(
         self.reason = reason
 
 
+class InvalidDirectiveStructure(Diagnostic):
+    severity = Diagnostic.Level.error
+
+    def __init__(
+        self,
+        msg: str,
+        start: Union[int, Tuple[int, int]],
+        end: Union[None, int, Tuple[int, int]] = None,
+    ) -> None:
+        super().__init__(f'Directive "io-code-block" {msg}', start, end)
+
+
 class InvalidInclude(Diagnostic):
     severity = Diagnostic.Level.error
 

snooty/parser.py

@@ -49,11 +49,13 @@
     ExpectedPathArg,
     FetchError,
     ImageSuggested,
+    InvalidDirectiveStructure,
     InvalidField,
     InvalidLiteralInclude,
     InvalidTableStructure,
     InvalidURL,
     MalformedGlossary,
+    MissingChild,
     RemovedLiteralBlockSyntax,
     TabMustBeDirective,
     TodoInfo,
@@ -454,6 +456,9 @@ def dispatch_departure(self, node: docutils.nodes.Node) -> None:
         elif isinstance(popped, n.Directive) and popped.name == "tabs":
             self.validate_tabs_children(popped)
 
+        elif isinstance(popped, n.Directive) and popped.name == "io-code-block":
+            self.diagnostics.extend(_validate_io_code_block_children(popped))
+
         elif (
             isinstance(popped, n.Directive)
             and f"{popped.domain}:{popped.name}" == ":glossary"
@@ -700,9 +705,12 @@ def create_page() -> Tuple[Page, EmbeddedRstParser]:
                 )
                 return doc
 
-        elif name == "literalinclude":
+        elif name == "literalinclude" or name == "input" or name == "output":
+            if name == "literalinclude":
+                if argument_text is None:
+                    self.diagnostics.append(ExpectedPathArg(name, line))
+                    return doc
             if argument_text is None:
-                self.diagnostics.append(ExpectedPathArg(name, line))
                 return doc
 
             _, filepath = util.reroot_path(
@@ -722,80 +730,88 @@ def create_page() -> Tuple[Page, EmbeddedRstParser]:
                 return doc
 
             lines = text.split("\n")
+            # Capture the original file-length before splicing it
+            len_file = len(lines)
 
-            def _locate_text(text: str) -> int:
-                """
-                Searches the literally-included file ('lines') for the specified text. If no such text is found,
-                add an InvalidLiteralInclude diagnostic.
-                """
-                assert isinstance(text, str)
-                loc = next((idx for idx, line in enumerate(lines) if text in line), -1)
-                if loc < 0:
-                    self.diagnostics.append(
-                        InvalidLiteralInclude(f'"{text}" not found in {filepath}', line)
-                    )
-                return loc
-
-            # Locate the start_after query
-            start_after = 0
-            if "start-after" in options:
-                start_after_text = options["start-after"]
-                start_after = _locate_text(start_after_text)
-                # Only increment start_after if text is specified, to avoid capturing the start_after_text
-                start_after += 1
-
-            # ...now locate the end_before query
-            end_before = len(lines)
-            if "end-before" in options:
-                end_before_text = options["end-before"]
-                end_before = _locate_text(end_before_text)
-
-            # Check that start_after_text precedes end_before_text (and end_before exists)
-            if start_after >= end_before >= 0:
-                self.diagnostics.append(
-                    InvalidLiteralInclude(
-                        f'"{end_before_text}" precedes "{start_after_text}" in {filepath}',
-                        line,
-                    )
-                )
+            if name == "literalinclude":
 
-            # If we failed to locate end_before text, default to the end-of-file
-            if end_before == -1:
+                def _locate_text(text: str) -> int:
+                    """
+                    Searches the literally-included file ('lines') for the specified text. If no such text is found,
+                    add an InvalidLiteralInclude diagnostic.
+                    """
+                    assert isinstance(text, str)
+                    loc = next(
+                        (idx for idx, line in enumerate(lines) if text in line), -1
+                    )
+                    if loc < 0:
+                        self.diagnostics.append(
+                            InvalidLiteralInclude(
+                                f'"{text}" not found in {filepath}', line
+                            )
+                        )
+                    return loc
+
+                # Locate the start_after query
+                start_after = 0
+                if "start-after" in options:
+                    start_after_text = options["start-after"]
+                    # start_after = self._locate_text(start_after_text, lines, line, text)
+                    start_after = _locate_text(start_after_text)
+                    # Only increment start_after if text is specified, to avoid capturing the start_after_text
+                    start_after += 1
+
+                # ...now locate the end_before query
                 end_before = len(lines)
+                if "end-before" in options:
+                    end_before_text = options["end-before"]
+                    # end_before = self._locate_text(end_before_text, lines, line, text)
+                    end_before = _locate_text(end_before_text)
 
-            # Capture the original file-length before splicing it
-            len_file = len(lines)
-            lines = lines[start_after:end_before]
-
-            dedent = 0
-            if "dedent" in options:
-                # Dedent is specified as a flag
-                if isinstance(options["dedent"], bool):
-                    # Deduce a reasonable dedent
-                    try:
-                        dedent = min(
-                            len(line) - len(line.lstrip())
-                            for line in lines
-                            if len(line.lstrip()) > 0
-                        )
-                    except ValueError:
-                        # Handle the (unlikely) case where there are no non-empty lines
-                        dedent = 0
-                # Dedent is specified as a nonnegative integer (number of characters):
-                # Note: since boolean is a subtype of int, this conditonal must follow the
-                # above bool-type conditional.
-                elif isinstance(options["dedent"], int):
-                    dedent = options["dedent"]
-                else:
+                # Check that start_after_text precedes end_before_text (and end_before exists)
+                if start_after >= end_before >= 0:
                     self.diagnostics.append(
                         InvalidLiteralInclude(
-                            f'Dedent "{dedent}" of type {type(dedent)}; expected nonnegative integer or flag',
+                            f'"{end_before_text}" precedes "{start_after_text}" in {filepath}',
                             line,
                         )
                     )
-                    return doc
 
-            lines = [line[dedent:] for line in lines]
+                # If we failed to locate end_before text, default to the end-of-file
+                if end_before == -1:
+                    end_before = len(lines)
+
+                lines = lines[start_after:end_before]
+
+                dedent = 0
+                if "dedent" in options:
+                    # Dedent is specified as a flag
+                    if isinstance(options["dedent"], bool):
+                        # Deduce a reasonable dedent
+                        try:
+                            dedent = min(
+                                len(line) - len(line.lstrip())
+                                for line in lines
+                                if len(line.lstrip()) > 0
+                            )
+                        except ValueError:
+                            # Handle the (unlikely) case where there are no non-empty lines
+                            dedent = 0
+                    # Dedent is specified as a nonnegative integer (number of characters):
+                    # Note: since boolean is a subtype of int, this conditonal must follow the
+                    # above bool-type conditional.
+                    elif isinstance(options["dedent"], int):
+                        dedent = options["dedent"]
+                    else:
+                        self.diagnostics.append(
+                            InvalidLiteralInclude(
+                                f'Dedent "{dedent}" of type {type(dedent)}; expected nonnegative integer or flag',
+                                line,
+                            )
+                        )
+                        return doc
+
+                lines = [line[dedent:] for line in lines]
 
             emphasize_lines = None
             if "emphasize-lines" in options:
@@ -806,7 +822,8 @@ def _locate_text(text: str) -> int:
                 except ValueError as err:
                     self.diagnostics.append(
                         InvalidLiteralInclude(
-                            f"Invalid emphasize-lines specification caused: {err}", line
+                            f"Invalid emphasize-lines specification caused: {err}",
+                            line,
                         )
                     )
 
@@ -833,6 +850,16 @@ def _locate_text(text: str) -> int:
 
             doc.children.append(code)
 
+        elif name == "io-code-block":
+            if argument_text is not None:
+                self.diagnostics.append(
+                    InvalidDirectiveStructure(
+                        "did not expect an argument, language should be passed as an option to input/output directives",
+                        line,
+                    )
+                )
+                return doc
+
         elif name == "include":
             if argument_text is None:
                 self.diagnostics.append(ExpectedPathArg(name, util.get_line(node)))
@@ -984,6 +1011,75 @@ def __make_child_visitor(self) -> "JSONVisitor":
         return visitor
 
 
+def _validate_io_code_block_children(node: n.Directive) -> List[Diagnostic]:
+    """Validates that a given io-code-block directive has 1 input and 1 output
+    child nodes, and copies the io-code-block's options into the options of the
+    underlying code nodes."""
+    # new_children should contain input and output directives
+    new_children: List[n.Node] = []
+    line = node.start[0]
+    expected_children = {"input", "output"}
+    diagnostics: List[Diagnostic] = []
+
+    for child in node.children:
+        if not isinstance(child, n.Directive):
+            diagnostics.append(
+                InvalidDirectiveStructure(
+                    f"expected input/output child directives, saw {child.type}",
+                    line,
+                )
+            )
+            continue
+
+        if child.name in expected_children:
+            new_grandchildren: List[n.Node] = []
+
+            # Input or output should have 1 child Code node
+            if len(child.children) == 1:
+                expected_children.remove(child.name)
+
+                # child nodes for input/output will inherit parent options
+                grandchild = child.children[0]
+                if isinstance(grandchild, n.Code):
+                    grandchild.lang = (
+                        child.options["language"]
+                        if "language" in child.options
+                        else None
+                    )
+                    grandchild.caption = (
+                        node.options["caption"] if "caption" in node.options else None
+                    )
+                    grandchild.copyable = (
+                        True
+                        if "copyable" in node.options and node.options["copyable"]
+                        else False
+                    )
+                    new_grandchildren.append(grandchild)
+                child.children = new_grandchildren
+                new_children.append(child)
+        else:
+            # either duplicate input/output or invalid child is provided
+            msg = f"already contains an {child.name} directive"
+            if not (child.name == "input" or child.name == "output"):
+                msg = f"does not accept child {child.name}"
+            diagnostics.append(
+                InvalidDirectiveStructure(
+                    msg,
+                    line,
+                )
+            )
+
+    # handle missing nested input and/or output directives
+    if len(expected_children) != 0 or len(new_children) != 2:
+        for expected_child in expected_children:
+            diagnostics.append(MissingChild("io-code-block", expected_child, line))
+            if expected_child == "input":
+                new_children = []
+
+    node.children = new_children
+    return diagnostics
+
+
 class InlineJSONVisitor(JSONVisitor):
     """A JSONVisitor subclass which does not emit block nodes."""
 

snooty/rstparser.py

@@ -809,6 +809,55 @@ def run(self) -> List[docutils.nodes.Node]:
         return [node]
 
 
+class BaseCodeIODirective(docutils.parsers.rst.Directive):
+    """Special handling for code input/output directives.
+
+    These directives can either take in a filepath or raw code content. If a filepath
+    is present, this should be included in the `argument` field of the AST. If raw code
+    content is present, it should become the value of the child Code node.
+    """
+
+    optional_arguments = 1
+
+    def run(self) -> List[docutils.nodes.Node]:
+        source, line = self.state_machine.get_source_and_line(self.lineno)
+        copyable = "copyable" not in self.options or self.options["copyable"]
+        linenos = "linenos" in self.options
+
+        node = directive("", self.name)
+        node.document = self.state.document
+        node.source, node.line = source, line
+        node["options"] = self.options
+
+        if self.arguments:
+            title_node: docutils.nodes.Node = docutils.nodes.Text(self.arguments[0])
+            node.append(directive_argument(self.arguments[0], "", title_node))
+        else:
+            try:
+                n_lines = len(self.content)
+                emphasize_lines = parse_linenos(
+                    self.options.get("emphasize-lines", ""), n_lines
+                )
+            except ValueError as err:
+                error_node = self.state.document.reporter.error(
+                    str(err), line=self.lineno
+                )
+                return [error_node]
+
+            value = "\n".join(self.content)
+            child_code = code(value, value)
+            child_code["name"] = "code"
+            child_code["emphasize_lines"] = emphasize_lines
+            child_code["linenos"] = linenos
+            child_code["copyable"] = copyable
+
+            child_code.document = self.state.document
+            child_code.source, node.line = source, line
+            node.append(child_code)
+
+        return [node]
+
+
 class BaseVersionDirective(docutils.parsers.rst.Directive):
     """Special handling for version change directives.
 
@@ -1030,6 +1079,8 @@ def get(cls, default_domain: Optional[str]) -> "Registry":
 SPECIAL_DIRECTIVE_HANDLERS: Dict[str, Type[docutils.parsers.rst.Directive]] = {
     "code-block": BaseCodeDirective,
     "code": BaseCodeDirective,
+    "input": BaseCodeIODirective,
+    "output": BaseCodeIODirective,
     "sourcecode": BaseCodeDirective,
     "versionadded": BaseVersionDirective,
     "versionchanged": BaseVersionDirective,