From ddc6b8bcd10e18ff844cbaeb777777d2612bcb4b Mon Sep 17 00:00:00 2001
From: Casper Weiss Bang <master@thecdk.net>
Date: Wed, 21 Oct 2020 20:16:53 +0200
Subject: [PATCH] Check both *.py and *.pyi files

---
 .gitignore          |  4 ++++
 pytest_black.py     |  2 +-
 tests/test_black.py | 27 +++++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index a18b0b5..8fcda9d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,7 @@
+# virtual environments
+
+.venv/
+.venv
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
diff --git a/pytest_black.py b/pytest_black.py
index 04c80cb..44a66c5 100644
--- a/pytest_black.py
+++ b/pytest_black.py
@@ -22,7 +22,7 @@ def pytest_addoption(parser):
 
 def pytest_collect_file(path, parent):
     config = parent.config
-    if config.option.black and path.ext == ".py":
+    if config.option.black and path.ext in [".py", ".pyi"]:
         if hasattr(BlackItem, "from_parent"):
             return BlackItem.from_parent(parent, fspath=path)
         else:
diff --git a/tests/test_black.py b/tests/test_black.py
index 0405169..f49ecf3 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 
 import pytest
+from py.path import local
+from pytest_black import pytest_collect_file
 
 
 pytestmark = pytest.mark.usefixtures('black_available')
@@ -176,3 +178,28 @@ def hello():
 
     out = "\n".join(result.stdout.lines)
     assert "PytestUnknownMarkWarning" not in out
+
+
+def test_gathers_pyi_files(tmpdir, request):
+    """Assert that pytest_collect_file handles *.pyi files"""
+    config = request.session
+    config.config.option.black = True
+    assert pytest_collect_file(local("test.pyi"), config) is not None
+
+
+def test_gathers_py_files(tmpdir, request):
+    """Assert that pytest_collect_file handles *.py files"""
+    config = request.session
+    config.config.option.black = True
+    assert pytest_collect_file(local("test.py"), config) is not None
+
+
+def test_ignores_pyc_files(tmpdir, request):
+    """
+    Assert that pytest_collect_file ignores *.pyc files
+
+    Used as an example of a non *.pyi and *.py)
+    """
+    config = request.session
+    config.config.option.black = True
+    assert pytest_collect_file(local("test.pyc"), config) is None