From da3449f9089b4742d3a516bbf2727b7ebc88d48b Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 23 Oct 2024 22:28:12 +0100
Subject: [PATCH] Add support for Black's `--pyi` flag (#375)

Fixes #374

---------

Co-authored-by: Adam Johnson <me@adamj.eu>
---
 CHANGELOG.rst                |  7 +++++++
 README.rst                   |  5 +++++
 src/blacken_docs/__init__.py |  8 +++++---
 tests/test_blacken_docs.py   | 28 ++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 7c3a26e..56c31b3 100644
--- 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
index 351a240..2ca70a9 100644
--- 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
index 1e1aa8c..247d19f 100644
--- 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
index 18bdbe4..b0d0e36 100644
--- 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(