updated haiku guide

docs/_ext/codediff.py

@@ -42,27 +42,45 @@ def parse(
     self,
     lines,
     title_left='Base',
+    title_middle=None,
     title_right='Diff',
     code_sep='---',
     sync=MISSING,
+    skip_left_doctest=MISSING,
   ):
     sync = sync is not MISSING
+    # skip legacy code snippets in upgrade guides
+    skip_left_doctest=skip_left_doctest is not MISSING
 
     if code_sep not in lines:
       raise ValueError(
         'Code separator not found! Code snippets should be '
         f'separated by {code_sep}.'
       )
-    idx = lines.index(code_sep)
-    code_left = self._code_block(lines[0:idx])
-    test_code = lines[idx + 1 :]
-    code_right = self._code_block(test_code)
-
-    output = self._tabs(
-      (title_left, code_left), (title_right, code_right), sync=sync
-    )
+    idxs = [i for i, x in enumerate(lines) if x == code_sep]
+    code_left = self._code_block(lines[0:idxs[0]])
+
+    test_codes = []
+    if not skip_left_doctest:
+      test_codes.extend(lines[0:idxs[0]])
+    if title_middle is None:
+      assert len(idxs) == 1
+      test_codes.extend(lines[idxs[0] + 1 :])
+      code_right = self._code_block(lines[idxs[0] + 1 :])
+      output = self._tabs(
+        (title_left, code_left), (title_right, code_right), sync=sync
+      )
+    else:
+      assert len(idxs) == 2
+      test_codes.extend(lines[idxs[0] + 1 : idxs[1]])
+      code_middle = self._code_block(lines[idxs[0] + 1 : idxs[1]])
+      test_codes.extend(lines[idxs[1] + 1 :])
+      code_right = self._code_block(lines[idxs[1] + 1 :])
+      output = self._tabs(
+        (title_left, code_left), (title_middle, code_middle), (title_right, code_right), sync=sync
+      )
 
-    return output, test_code
+    return output, test_codes
 
   def _code_block(self, lines):
     """Creates a codeblock."""
@@ -100,9 +118,11 @@ class CodeDiffDirective(SphinxDirective):
   has_content = True
   option_spec = {
     'title_left': directives.unchanged,
+    'title_middle': directives.unchanged,
     'title_right': directives.unchanged,
     'code_sep': directives.unchanged,
     'sync': directives.flag,
+    'skip_left_doctest': directives.flag,
   }
 
   def run(self):
@@ -114,7 +134,7 @@ def run(self):
     # We add attribute "testnodetype" so it is be picked up by the doctest
     # builder. This functionality is not officially documented but can be found
     # in the source code:
-    # https://github.com/sphinx-doc/sphinx/blob/3.x/sphinx/ext/doctest.py
+    # https://github.com/sphinx-doc/sphinx/blob/master/sphinx/ext/doctest.py
     # (search for 'testnodetype').
     test_code = '\n'.join(test_code)
     test_node = nodes.comment(test_code, test_code, testnodetype='testcode')