Merge pull request #169 from snejus/fix-prompt-continuation-lines

Fix formatting Python code blocks with continuation lines

Author: LilSpazJoekp

CHANGES.rst

@@ -33,6 +33,7 @@
 - Image and Unicode substitution definitions were not being formatted correctly.
 - Files specified in pyproject are now parsed correctly.
 - Fixed issue when loading cache in certain environments.
+- Fixed formatting of Python code blocks with continuation lines in docstrings.
 
 **Removed**
 

docstrfmt/docstrfmt.py

@@ -1075,20 +1075,10 @@ def doctest_block(
                 f"Invalid doctest block: {e}",
             ) from None
         doctest_blocks = []
-        current_block = ""
-        while parsed:
-            example = parsed.pop(0)
-            if example.want:
-                current_block += example.source.strip()
-                formatted = CodeFormatters(current_block, context).python()
-                doctest_blocks.append((formatted, example.want.strip()))
-                current_block = ""
-            else:  # pragma: no cover
-                current_block += example.source.strip() + "\n"
-        if current_block:
-            # If there's any remaining code, format it as well.
-            formatted = CodeFormatters(current_block, context).python()
-            doctest_blocks.append((formatted, ""))
+        for example in parsed:
+            formatted = CodeFormatters(example.source.strip(), context).python()
+            doctest_blocks.append((formatted, example.want.strip()))
+
         if not doctest_blocks:
             raise InvalidRstError(
                 context.current_file,
@@ -1097,11 +1087,10 @@ def doctest_block(
                 "Empty doctest block.",
             )
         for formatted, want in doctest_blocks:
-            for line in formatted.splitlines():
-                if line.startswith(" "):
-                    yield f"... {line}"
-                elif line:
-                    yield f">>> {line}"
+            first_line, *other_lines = formatted.splitlines()
+            yield f">>> {first_line}"
+            for line in other_lines:
+                yield f"... {line}" if line else "..."
             if want:
                 yield from want.splitlines()
 

tests/test_files/py_file.py

@@ -122,6 +122,16 @@ def method2(self, attr1, attr2):
         4
         >>> import this
 
+        >>> try:
+        ...     value = {1: 2}[3]
+        ... except KeyError:
+        ...     value = 4
+        ...
+        ... print(
+        ...     {"some long long long long long long long long long key to force line wrap": value}
+        ... )
+        {"some long long long long long long long long long key to force line wrap": value}
+
         """
         my_position, im_active = 1, True
         match my_position, im_active: