diff --git a/pyproject.toml b/pyproject.toml
index 6ca4f147..3070484d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,6 +36,7 @@ requests = { version ="^2.22.0", optional = true }
 pytest = { version= "^8.0.0", optional = true }
 pytest-benchmark = { version = "^4.0", optional = true }
 flake8 = { version = "^5.0.4", optional = true }
+black = { version = ">=24.8.0", optional = true }
 
 [tool.poetry.extras]
 dev = [
@@ -44,8 +45,12 @@ dev = [
     "pytest",
     "pytest-benchmark",
     "requests",
+    "black",
 ]
 
+[tool.black]
+include = "(^datadog_lambda/.*\\.pyi?$|^tests/.*\\.pyi?$)"
+
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
diff --git a/scripts/check_format.sh b/scripts/check_format.sh
index 783a17d6..f630a425 100755
--- a/scripts/check_format.sh
+++ b/scripts/check_format.sh
@@ -1,9 +1,6 @@
 #!/bin/sh
 set -e
 
-PYTHON_VERSION=$(python -c 'import sys; print(sys.version_info.major)')
-pip install -Iv black==25.11.0
-
 python -m black --check datadog_lambda/ --diff
 python -m black --check tests --diff