fix: format code snippets properly (#193)

* fix: format code snippets properly

* test: update unit test

* test: add unit test for testing with blocks of code snippet

* test: parameterize unit test for indent_left

* docs: add docstring to indent_code_left method

Author: dandhlee

docfx_yaml/extension.py

@@ -386,13 +386,17 @@ def extract_keyword(line):
         return line
 
 
-# Given lines of code, indent to left by 1 block, based on
-# amount of trailing white space of first line as 1 block.
-def indent_code_left(lines):
+def indent_code_left(lines, tab_space):
+    """Indents code lines left by tab_space.
+
+    Args:
+        lines: String lines of code.
+        tab_space: Number of spaces to indent to left by.
+
+    Returns:
+        String lines of left-indented code.
+    """
     parts = lines.split("\n")
-    # Count how much leading whitespaces there are based on first line.
-    # lstrip(" ") removes all trailing whitespace from the string.
-    tab_space = len(parts[0]) - len(parts[0].lstrip(" "))
     parts = [part[tab_space:] for part in parts]
     return "\n".join(parts)
 
@@ -417,6 +421,9 @@ def _parse_docstring_summary(summary):
         if keyword and keyword in [CODE, CODEBLOCK]:
             # If we reach the end of keyword, close up the code block.
             if not part.startswith(" "*tab_space) or part.startswith(".."):
+                if code_snippet:
+                    parts = [indent_code_left(part, tab_space) for part in code_snippet]
+                    summary_parts.append("\n\n".join(parts))
                 summary_parts.append("```\n")
                 keyword = ""
 
@@ -428,9 +435,12 @@ def _parse_docstring_summary(summary):
                         raise ValueError(f"Code in the code block should be indented. Please check the docstring: \n{summary}")
                 if not part.startswith(" "*tab_space):
                     # No longer looking at code-block, reset keyword.
+                    if code_snippet:
+                        parts = [indent_code_left(part, tab_space) for part in code_snippet]
+                        summary_parts.append("\n\n".join(parts))
                     keyword = ""
                     summary_parts.append("```\n")
-                summary_parts.append(indent_code_left(part))
+                code_snippet.append(part)
                 continue
 
         # Attributes come in 3 parts, parse the latter two here.
@@ -472,6 +482,7 @@ def _parse_docstring_summary(summary):
                 language = part.split("::")[1].strip()
                 summary_parts.append(f"```{language}")
 
+                code_snippet = []
                 tab_space = -1
 
             # Extract the name for attribute first.
@@ -490,6 +501,9 @@ def _parse_docstring_summary(summary):
     # Close up from the keyword if needed.
     if keyword and keyword in [CODE, CODEBLOCK]:
         # Check if it's already closed.
+        if code_snippet:
+            parts = [indent_code_left(part, tab_space) for part in code_snippet]
+            summary_parts.append("\n\n".join(parts))
         if summary_parts[-1] != "```\n":
             summary_parts.append("```\n")
 

tests/test_helpers.py

@@ -14,24 +14,33 @@
 import tempfile
 
 class TestGenerate(unittest.TestCase):
-    def test_indent_code_left(self):
+    code_testdata = [
         # Check that the code indents to left based on first line.
-        code_want = \
-"""def foo():
-    print('test function for indent')
-    return ('left-indented-code')
-"""
-
-        code = \
+        [
+            \
 """    def foo():
         print('test function for indent')
         return ('left-indented-code')
+""",
+            \
+"""def foo():
+    print('test function for indent')
+    return ('left-indented-code')
 """
-        code = indent_code_left(code)
-        self.assertEqual(code, code_want)
-
+        ],
         # Check that if there's no whitespace, it does not indent
-        code_want = \
+        [
+            \
+"""
+print('test function for no impact indent')
+for i in range(10):
+    print(i)
+    if i%5 == 0:
+        i += 1
+    else:
+        continue
+""",
+            \
 """
 print('test function for no impact indent')
 for i in range(10):
@@ -41,9 +50,34 @@ def test_indent_code_left(self):
     else:
         continue
 """
+        ],
+    ]
+    @parameterized.expand(code_testdata)
+    def test_indent_code_left(self, code, code_want):
+        parts = code.split("\n")
+        tab_space = len(parts[0]) - len(parts[0].lstrip(" "))
+        code_got = indent_code_left(code, tab_space)
+        self.assertEqual(code_got, code_want)
+
 
-        code_got = indent_code_left(code_want)
-        # Confirm that nothing changes.
+    def test_indent_code_blocks_left(self):
+        # Check code blocks are indented properly.
+        code_want = \
+"""def foo():
+
+    print('test function for indent')
+
+    return ('left-indented-blocks')
+"""
+
+        # Test with how blocks would appear in the code block
+        code = [
+            "    def foo():",
+            "        print('test function for indent')",
+            "        return ('left-indented-blocks')\n"
+        ]
+        tab_space = len(code[0]) - len(code[0].lstrip(" "))
+        code_got = "\n\n".join([indent_code_left(part, tab_space) for part in code])
         self.assertEqual(code_got, code_want)