Check both *.py and *.pyi files

.gitignore

@@ -1,3 +1,6 @@
+# virtual environments
+
+.venv/
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

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:

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