Add support for Black's --pyi flag (#375)

Fixes #374 --------- Co-authored-by: Adam Johnson <[email protected]>
adamchainz · Oct 23, 2024 · da3449f · da3449f
1 parent a05853b
commit da3449f
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 3 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -2,6 +2,13 @@
 Changelog
 =========
 
+Unreleased
+----------
+
+* Add support for Black's ``--pyi`` flag.
+
+  Thanks to Alex Waygood in `PR #375 <https://github.com/adamchainz/blacken-docs/pull/375>`__.
+
 1.19.0 (2024-10-07)
 -------------------
 

diff --git a/README.rst b/README.rst
@@ -104,6 +104,11 @@ blacken-docs currently passes the following options through to Black:
   .. |--preview| replace:: ``--preview``
   __ https://black.readthedocs.io/en/stable/usage_and_configuration/the_basics.html#preview
 
+* |--pyi|__
+
+  .. |--pyi| replace:: ``--pyi``
+  __ https://black.readthedocs.io/en/stable/usage_and_configuration/the_basics.html#pyi
+
 * |-S / --skip-string-normalization|__
 
   .. |-S / --skip-string-normalization| replace:: ``-S`` / ``--skip-string-normalization``

diff --git a/src/blacken_docs/__init__.py b/src/blacken_docs/__init__.py
@@ -338,14 +338,16 @@ def main(argv: Sequence[str] | None = None) -> int:
         "--rst-literal-blocks",
         action="store_true",
     )
+    parser.add_argument("--pyi", action="store_true")
     parser.add_argument("filenames", nargs="*")
     args = parser.parse_args(argv)
 
-    black_mode = black.FileMode(
+    black_mode = black.Mode(
+        target_versions=set(args.target_versions),
         line_length=args.line_length,
-        preview=args.preview,
         string_normalization=not args.skip_string_normalization,
-        target_versions=set(args.target_versions),
+        is_pyi=args.pyi,
+        preview=args.preview,
     )
 
     retv = 0

diff --git a/tests/test_blacken_docs.py b/tests/test_blacken_docs.py
@@ -866,6 +866,34 @@ def test_integration_preview(tmp_path):
     )
 
 
+def test_integration_pyi(tmp_path):
+    f = tmp_path / "f.md"
+    f.write_text(
+        dedent(
+            """\
+            ```python
+            class Foo: ...
+
+
+            class Bar: ...
+            ```
+            """
+        )
+    )
+
+    result = blacken_docs.main((str(f), "--pyi"))
+
+    assert result == 1
+    assert f.read_text() == dedent(
+        """\
+        ```python
+        class Foo: ...
+        class Bar: ...
+        ```
+        """
+    )
+
+
 def test_integration_py36(tmp_path):
     f = tmp_path / "f.md"
     f.write_text(