Run ruff on entire project

Makefile

@@ -33,22 +33,22 @@ test/integration:
 
 .PHONY: lint
 lint: ## Lint project.
-	@poetry run ruff check --fix griptape/
+	@poetry run ruff check --fix
 
 .PHONY: format
 format: ## Format project.
-	@poetry run ruff format .
+	@poetry run ruff format
 
 .PHONY: check
 check: check/format check/lint check/types check/spell ## Run all checks.
 
 .PHONY: check/format
 check/format:
-	@poetry run ruff format --check griptape/
+	@poetry run ruff format --check
 
 .PHONY: check/lint
 check/lint:
-	@poetry run ruff check griptape/
+	@poetry run ruff check
 
 .PHONY: check/types
 check/types:

docs/gen_ref_pages.py

@@ -1,11 +1,12 @@
 """Generate the code reference pages and navigation."""
 
-from textwrap import dedent
 from pathlib import Path
+from textwrap import dedent
+
 import mkdocs_gen_files
 
 
-def build_reference_docs():
+def build_reference_docs() -> None:
     nav = mkdocs_gen_files.Nav()
 
     for path in sorted(Path("griptape").rglob("*.py")):
@@ -37,8 +38,8 @@ def build_reference_docs():
         index_file.write(
             dedent(
                 """
-        # Overview 
-        This section of the documentation is dedicated to a reference API of Griptape. 
+        # Overview
+        This section of the documentation is dedicated to a reference API of Griptape.
         Here you will find every class, function, and method that is available to you when using the library.
         """
             )

docs/plugins/swagger_ui_plugin.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any
 
 import markdown
 from jinja2 import Environment, FileSystemLoader, select_autoescape
@@ -11,7 +12,7 @@
 }
 
 
-def generate_page_contents(page):
+def generate_page_contents(page: Any) -> str:
     spec_url = config_scheme["spec_url"]
     tmpl_url = config_scheme["template"]
     env = Environment(loader=FileSystemLoader("docs/plugins/tmpl"), autoescape=select_autoescape(["html", "xml"]))
@@ -23,11 +24,11 @@ def generate_page_contents(page):
     return tmpl_out
 
 
-def on_config(config):
-    print("INFO     -  swagger-ui plugin ENABLED")
+def on_config(config: Any) -> None:
+    pass
 
 
-def on_page_read_source(page, config):
+def on_page_read_source(page: Any, config: Any) -> Any:
     index_path = os.path.join(config["docs_dir"], config_scheme["outfile"])
     page_path = os.path.join(config["docs_dir"], page.file.src_path)
     if index_path == page_path:

pyproject.toml

@@ -217,6 +217,7 @@ select = [
     "C4", # flake8-comprehensions
     "ANN", # flake8-annotations
     "FBT", # flake8-boolean-trap
+    "PT", # flake8-pytest-style
 ]
 ignore = [
     "UP007", # non-pep604-annotation
@@ -238,12 +239,20 @@ ignore = [
     "ANN101", # missing-type-self
     "ANN102", # missing-type-cls
     "ANN401", # any-type
+    "PT011", # pytest-raises-too-broad
 ]
 [tool.ruff.lint.pydocstyle]
 convention = "google"
 
 [tool.ruff.lint.per-file-ignores]
-"__init__.py" = ["I"]
+"__init__.py" = [
+    "I" # isort
+]
+"tests/*" = [
+    "ANN001", # missing-type-function-argument
+    "ANN201", # missing-return-type-undocumented-public-function
+    "ANN202", # missing-return-type-private-function
+]
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "attr".msg = "The attr module is deprecated, use attrs instead."

tests/integration/drivers/vector/test_pgvector_vector_store_driver.py

@@ -1,9 +1,11 @@
 import uuid
+
 import pytest
+from sqlalchemy import create_engine
+
 from griptape.drivers import PgVectorVectorStoreDriver
 from tests.mocks.mock_embedding_driver import MockEmbeddingDriver
 from tests.utils.postgres import can_connect_to_postgres
-from sqlalchemy import create_engine
 
 
 @pytest.mark.skipif(not can_connect_to_postgres(), reason="Postgres is not present")
@@ -13,11 +15,11 @@ class TestPgVectorVectorStoreDriver:
     vec1 = [0.1, 0.2, 0.3]
     vec2 = [0.4, 0.5, 0.6]
 
-    @pytest.fixture
+    @pytest.fixture()
     def embedding_driver(self):
         return MockEmbeddingDriver()
 
-    @pytest.fixture
+    @pytest.fixture()
     def vector_store_driver(self, embedding_driver):
         driver = PgVectorVectorStoreDriver(
             connection_string=self.connection_string, embedding_driver=embedding_driver, table_name=self.table_name
@@ -28,13 +30,13 @@ def vector_store_driver(self, embedding_driver):
         return driver
 
     def test_initialize_requires_engine_or_connection_string(self, embedding_driver):
+        driver = PgVectorVectorStoreDriver(embedding_driver=embedding_driver, table_name=self.table_name)
         with pytest.raises(ValueError):
-            driver = PgVectorVectorStoreDriver(embedding_driver=embedding_driver, table_name=self.table_name)
             driver.setup()
 
     def test_initialize_accepts_engine(self, embedding_driver):
         engine = create_engine(self.connection_string)
-        driver = PgVectorVectorStoreDriver(embedding_driver=embedding_driver, engine=engine, table_name=self.table_name)
+        driver = PgVectorVectorStoreDriver(embedding_driver=embedding_driver, engine=engine, table_name=self.table_name)  # pyright: ignore[reportArgumentType]
 
         driver.setup()
 
@@ -86,11 +88,9 @@ def test_can_insert_and_load_entry_with_namespace(self, vector_store_driver):
         assert result.vector == pytest.approx(self.vec1)
 
     def test_can_load_entries(self, vector_store_driver):
-        """
-        Depending on when this test is executed relative to the others,
-        we don't know exactly how many vectors will be returned. We can
-        ensure that at least two exist and confirm that those are found.
-        """
+        # Depending on when this test is executed relative to the others,
+        # we don't know exactly how many vectors will be returned. We can
+        # ensure that at least two exist and confirm that those are found.
         vec1_id = vector_store_driver.upsert_vector(self.vec1)
         vec2_id = vector_store_driver.upsert_vector(self.vec2)
 
@@ -173,8 +173,9 @@ def test_query_returns_vectors_when_requested(self, vector_store_driver):
         assert results[0].vector == pytest.approx(embedding)
 
     def test_can_use_custom_table_name(self, embedding_driver, vector_store_driver):
-        """This test ensures at least one row exists in the default table before specifying
-        a custom table name. After inserting another row, we should be able to query only one
+        """This test ensures at least one row exists in the default table before specifying a custom table name.
+
+        After inserting another row, we should be able to query only one
         vector from the table, and it should be the vector added to the table with the new name.
         """
         vector_store_driver.upsert_vector(self.vec1)

tests/integration/rules/test_rule.py

@@ -1,14 +1,15 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestRule:
     @pytest.fixture(
         autouse=True, params=StructureTester.RULE_CAPABLE_PROMPT_DRIVERS, ids=StructureTester.prompt_driver_id_fn
     )
     def structure_tester(self, request):
-        from griptape.structures import Agent
         from griptape.rules import Rule
+        from griptape.structures import Agent
 
         agent = Agent(prompt_driver=request.param, rules=[Rule("Your name is Tony.")])
 

tests/integration/tasks/test_csv_extraction_task.py

@@ -1,6 +1,7 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestCsvExtractionTask:
     @pytest.fixture(
@@ -9,9 +10,9 @@ class TestCsvExtractionTask:
         ids=StructureTester.prompt_driver_id_fn,
     )
     def structure_tester(self, request):
-        from griptape.tasks import ExtractionTask
-        from griptape.structures import Agent
         from griptape.engines import CsvExtractionEngine
+        from griptape.structures import Agent
+        from griptape.tasks import ExtractionTask
 
         columns = ["Name", "Age", "Address"]
 

tests/integration/tasks/test_json_extraction_task.py

@@ -1,6 +1,7 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestJsonExtractionTask:
     @pytest.fixture(
@@ -9,11 +10,12 @@ class TestJsonExtractionTask:
         ids=StructureTester.prompt_driver_id_fn,
     )
     def structure_tester(self, request):
-        from griptape.tasks import ExtractionTask
-        from griptape.structures import Agent
-        from griptape.engines import JsonExtractionEngine
         from schema import Schema
 
+        from griptape.engines import JsonExtractionEngine
+        from griptape.structures import Agent
+        from griptape.tasks import ExtractionTask
+
         # Define some JSON data
         user_schema = Schema({"users": [{"name": str, "age": int, "location": str}]}).json_schema("UserSchema")
 

tests/integration/tasks/test_prompt_task.py

@@ -1,6 +1,7 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestPromptTask:
     @pytest.fixture(

tests/integration/tasks/test_rag_task.py

@@ -1,7 +1,8 @@
+import pytest
+
 from tests.mocks.mock_prompt_driver import MockPromptDriver
 from tests.utils.defaults import rag_engine
 from tests.utils.structure_tester import StructureTester
-import pytest
 
 
 class TestRagTask:
@@ -11,10 +12,10 @@ class TestRagTask:
         ids=StructureTester.prompt_driver_id_fn,
     )
     def structure_tester(self, request):
+        from griptape.artifacts import TextArtifact
+        from griptape.drivers import LocalVectorStoreDriver, OpenAiEmbeddingDriver
         from griptape.structures import Agent
         from griptape.tasks import RagTask
-        from griptape.drivers import LocalVectorStoreDriver, OpenAiEmbeddingDriver
-        from griptape.artifacts import TextArtifact
 
         vector_store_driver = LocalVectorStoreDriver(embedding_driver=OpenAiEmbeddingDriver())
         artifact = TextArtifact("John Doe works as as software engineer at Griptape.")

tests/integration/tasks/test_text_summary_task.py

@@ -1,6 +1,7 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestTextSummaryTask:
     @pytest.fixture(
@@ -10,8 +11,8 @@ class TestTextSummaryTask:
     )
     def structure_tester(self, request):
         from griptape.engines.summary.prompt_summary_engine import PromptSummaryEngine
-        from griptape.tasks import TextSummaryTask
         from griptape.structures import Agent
+        from griptape.tasks import TextSummaryTask
 
         agent = Agent(conversation_memory=None, prompt_driver=request.param)
         agent.add_task(TextSummaryTask(summary_engine=PromptSummaryEngine(prompt_driver=request.param)))
@@ -21,17 +22,17 @@ def structure_tester(self, request):
     def test_summary_task(self, structure_tester):
         structure_tester.run(
             """
-                Meeting transcriot: 
-                Miguel: Hi Brant, I want to discuss the workstream  for our new product launch 
-                Brant: Sure Miguel, is there anything in particular you want to discuss? 
-                Miguel: Yes, I want to talk about how users enter into the product. 
-                Brant: Ok, in that case let me add in Namita. 
-                Namita: Hey everyone 
-                Brant: Hi Namita, Miguel wants to discuss how users enter into the product. 
-                Miguel: its too complicated and we should remove friction.  for example, why do I need to fill out additional forms?  I also find it difficult to find where to access the product when I first land on the landing page. 
-                Brant: I would also add that I think there are too many steps. 
-                Namita: Ok, I can work on the landing page to make the product more discoverable but brant can you work on the additional forms? 
-                Brant: Yes but I would need to work with James from another team as he needs to unblock the sign up workflow.  Miguel can you document any other concerns so that I can discuss with James only once? 
-                Miguel: Sure. 
+                Meeting transcriot:
+                Miguel: Hi Brant, I want to discuss the workstream  for our new product launch
+                Brant: Sure Miguel, is there anything in particular you want to discuss?
+                Miguel: Yes, I want to talk about how users enter into the product.
+                Brant: Ok, in that case let me add in Namita.
+                Namita: Hey everyone
+                Brant: Hi Namita, Miguel wants to discuss how users enter into the product.
+                Miguel: its too complicated and we should remove friction.  for example, why do I need to fill out additional forms?  I also find it difficult to find where to access the product when I first land on the landing page.
+                Brant: I would also add that I think there are too many steps.
+                Namita: Ok, I can work on the landing page to make the product more discoverable but brant can you work on the additional forms?
+                Brant: Yes but I would need to work with James from another team as he needs to unblock the sign up workflow.  Miguel can you document any other concerns so that I can discuss with James only once?
+                Miguel: Sure.
         """
         )

tests/integration/tasks/test_tool_task.py

@@ -1,6 +1,7 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestToolTask:
     @pytest.fixture(

tests/integration/tasks/test_toolkit_task.py

@@ -1,6 +1,7 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestToolkitTask:
     @pytest.fixture(
@@ -10,9 +11,10 @@ class TestToolkitTask:
     )
     def structure_tester(self, request):
         import os
-        from griptape.structures import Agent
-        from griptape.tools import WebScraper, WebSearch, TaskMemoryClient
+
         from griptape.drivers import GoogleWebSearchDriver
+        from griptape.structures import Agent
+        from griptape.tools import TaskMemoryClient, WebScraper, WebSearch
 
         return StructureTester(
             Agent(

tests/integration/test_code_blocks.py

@@ -2,8 +2,8 @@
 import os
 
 import pytest
-from tests.utils.code_blocks import get_all_code_blocks, check_py_string
 
+from tests.utils.code_blocks import check_py_string, get_all_code_blocks
 
 if "DOCS_ALL_CHANGED_FILES" in os.environ and os.environ["DOCS_ALL_CHANGED_FILES"] != "":
     docs_all_changed_files = os.environ["DOCS_ALL_CHANGED_FILES"].split()

tests/integration/tools/test_calculator.py

@@ -1,6 +1,7 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestCalculator:
     @pytest.fixture(

tests/integration/tools/test_file_manager.py

@@ -1,6 +1,7 @@
-from tests.utils.structure_tester import StructureTester
 import pytest
 
+from tests.utils.structure_tester import StructureTester
+
 
 class TestFileManager:
     @pytest.fixture(

tests/integration/tools/test_google_docs_client.py

@@ -1,5 +1,7 @@
-import pytest
 import os
+
+import pytest
+
 from tests.utils.structure_tester import StructureTester