From 6734557f7d0e6761bfe406cc54f328748d62ae05 Mon Sep 17 00:00:00 2001
From: Bryan Forbes <bryan@reigndropsfall.net>
Date: Thu, 29 Apr 2021 16:03:22 -0500
Subject: [PATCH] Set `is_pyi` if `stdin_filename` ends with `.pyi`

Fixes #2167
---
 src/black/__init__.py |  2 ++
 tests/test_black.py   | 28 +++++++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/black/__init__.py b/src/black/__init__.py
index 1c69cc41cdc..49d088b531d 100644
--- a/src/black/__init__.py
+++ b/src/black/__init__.py
@@ -755,6 +755,8 @@ def reformat_one(
             is_stdin = False
 
         if is_stdin:
+            if src.suffix == ".pyi":
+                mode = replace(mode, is_pyi=True)
             if format_stdin_to_stdout(fast=fast, write_back=write_back, mode=mode):
                 changed = Changed.YES
         else:
diff --git a/tests/test_black.py b/tests/test_black.py
index c643f27cacf..a6725715d3f 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -1578,7 +1578,33 @@ def test_reformat_one_with_stdin_filename(self) -> None:
                 mode=DEFAULT_MODE,
                 report=report,
             )
-            fsts.assert_called_once()
+            fsts.assert_called_once_with(
+                fast=True, write_back=black.WriteBack.YES, mode=DEFAULT_MODE
+            )
+            # __BLACK_STDIN_FILENAME__ should have been striped
+            report.done.assert_called_with(expected, black.Changed.YES)
+
+    def test_reformat_one_with_stdin_filename_pyi(self) -> None:
+        with patch(
+            "black.format_stdin_to_stdout",
+            return_value=lambda *args, **kwargs: black.Changed.YES,
+        ) as fsts:
+            report = MagicMock()
+            p = "foo.pyi"
+            path = Path(f"__BLACK_STDIN_FILENAME__{p}")
+            expected = Path(p)
+            black.reformat_one(
+                path,
+                fast=True,
+                write_back=black.WriteBack.YES,
+                mode=DEFAULT_MODE,
+                report=report,
+            )
+            fsts.assert_called_once_with(
+                fast=True,
+                write_back=black.WriteBack.YES,
+                mode=replace(DEFAULT_MODE, is_pyi=True),
+            )
             # __BLACK_STDIN_FILENAME__ should have been striped
             report.done.assert_called_with(expected, black.Changed.YES)