MAINT Enable ruff rules: DTZ, N, T10, TID, YTT

pyproject.toml

@@ -257,14 +257,19 @@ select = [
     "CPY001",  # missing-copyright-notice
     "D",  # https://docs.astral.sh/ruff/rules/#pydocstyle-d
     "DOC",  # https://docs.astral.sh/ruff/rules/#pydoclint-doc
+    "DTZ",  # https://docs.astral.sh/ruff/rules/#flake8-datetimez-dtz
     "F401",  # unused-import
     "I",  # isort
+    "N",  # https://docs.astral.sh/ruff/rules/#pep8-naming-n
     "PIE",  # https://docs.astral.sh/ruff/rules/#flake8-pie-pie
     "RET",  # https://docs.astral.sh/ruff/rules/#flake8-return-ret
     "SIM",  # https://docs.astral.sh/ruff/rules/#flake8-simplify-sim
+    "T10",  # https://docs.astral.sh/ruff/rules/#flake8-debugger-t10
     "TCH",  # https://docs.astral.sh/ruff/rules/#flake8-type-checking-tch
+    "TID",  # https://docs.astral.sh/ruff/rules/#flake8-tidy-imports-tid
     "UP",  # https://docs.astral.sh/ruff/rules/#pyupgrade-up
     "W",  # https://docs.astral.sh/ruff/rules/#pycodestyle-w
+    "YTT",  # https://docs.astral.sh/ruff/rules/#flake8-2020-ytt
 ]
 ignore = [
     "B903",  # class-as-data-structure (test helper classes use @apply_defaults pattern)
@@ -300,7 +305,7 @@ notice-rgx = "Copyright \\(c\\) Microsoft Corporation\\.\\s*\\n.*Licensed under
 # Temporary ignores for pyrit/ subdirectories until issue #1176
 # https://github.com/Azure/PyRIT/issues/1176 is fully resolved
 # TODO: Remove these ignores once the issues are fixed
-"pyrit/{auxiliary_attacks,ui}/**/*.py" = ["B905", "D101", "D102", "D103", "D104", "D105", "D106", "D107", "D401", "D404", "D417", "D418", "DOC102", "DOC201", "DOC202", "DOC402", "DOC501", "SIM101", "SIM108"]
+"pyrit/{auxiliary_attacks,ui}/**/*.py" = ["B905", "D101", "D102", "D103", "D104", "D105", "D106", "D107", "D401", "D404", "D417", "D418", "DOC102", "DOC201", "DOC202", "DOC402", "DOC501", "N", "SIM101", "SIM108"]
 # Backend API routes raise HTTPException handled by FastAPI, not true exceptions
 "pyrit/backend/**/*.py" = ["DOC501", "B008"]
 "pyrit/__init__.py" = ["D104"]

pyrit/auth/azure_storage_auth.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from urllib.parse import urlparse
 
 from azure.identity.aio import DefaultAzureCredential
@@ -31,7 +31,7 @@ async def get_user_delegation_key(blob_service_client: BlobServiceClient) -> Use
         Returns:
             UserDelegationKey: A user delegation key valid for one day.
         """
-        delegation_key_start_time = datetime.now()
+        delegation_key_start_time = datetime.now(tz=timezone.utc)
         delegation_key_expiry_time = delegation_key_start_time + timedelta(days=1)
 
         return await blob_service_client.get_user_delegation_key(
@@ -79,7 +79,7 @@ async def get_sas_token(container_url: str) -> str:
                 storage_account_name = parsed_url.netloc.split(".")[0]
 
                 # Set start_time 5 minutes before the current time to account for any clock skew
-                start_time = datetime.now() - timedelta(minutes=5)
+                start_time = datetime.now(tz=timezone.utc) - timedelta(minutes=5)
                 expiry_time = start_time + timedelta(days=1)
 
                 sas_token = generate_container_sas(

pyrit/cli/frontend_core.py

@@ -32,7 +32,7 @@
     HAS_TERMCOLOR = False
 
     # Create a dummy termcolor module for fallback
-    class termcolor:  # type: ignore
+    class termcolor:  # type: ignore  # noqa: N801
         """Dummy termcolor fallback for colored printing if termcolor is not installed."""
 
         @staticmethod
@@ -718,8 +718,8 @@ def get_default_initializer_discovery_path() -> Path:
     Returns:
         Path to the scenarios initializers directory.
     """
-    PYRIT_PATH = Path(__file__).parent.parent.resolve()
-    return PYRIT_PATH / "setup" / "initializers" / "scenarios"
+    pyrit_path = Path(__file__).parent.parent.resolve()
+    return pyrit_path / "setup" / "initializers" / "scenarios"
 
 
 async def print_scenarios_list_async(*, context: FrontendCore) -> int:

pyrit/cli/pyrit_shell.py

@@ -411,7 +411,7 @@ def do_clear(self, arg: str) -> None:
     # Shortcuts and aliases
     do_quit = do_exit
     do_q = do_exit
-    do_EOF = do_exit  # Ctrl+D on Unix, Ctrl+Z on Windows
+    do_EOF = do_exit  # Ctrl+D on Unix, Ctrl+Z on Windows  # noqa: N815
 
     def emptyline(self) -> bool:
         """

pyrit/datasets/seed_datasets/local/local_dataset_loader.py

@@ -92,7 +92,7 @@ def _register_local_datasets() -> None:
                     # We override __init__ to pass the specific file_path
 
                     def make_init(path: Path) -> Callable[[Any], None]:
-                        def __init__(self: Any) -> None:
+                        def __init__(self: Any) -> None:  # noqa: N807
                             super(self.__class__, self).__init__(file_path=path)
 
                         return __init__

pyrit/exceptions/exception_classes.py

@@ -113,7 +113,7 @@ def __call__(self, retry_state: RetryCallState) -> float:
         return wait_strategy(retry_state)
 
 
-class PyritException(Exception, ABC):
+class PyritException(Exception, ABC):  # noqa: N818
     """Base exception class for PyRIT components."""
 
     def __init__(self, *, status_code: int = 500, message: str = "An error occurred") -> None:

pyrit/executor/attack/printer/console_printer.py

@@ -3,7 +3,7 @@
 
 import json
 import textwrap
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any
 
 from colorama import Back, Fore, Style
@@ -345,10 +345,10 @@ def _print_footer(self) -> None:
 
         Displays the current timestamp when the report was generated.
         """
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        timestamp = datetime.now(tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
         print()
         self._print_colored("─" * self._width, Style.DIM, Fore.WHITE)
-        footer_text = f"Report generated at: {timestamp}"
+        footer_text = f"Report generated at: {timestamp} UTC"
         self._print_colored(footer_text.center(self._width), Style.DIM, Fore.WHITE)
 
     def _print_section_header(self, title: str) -> None:

pyrit/executor/attack/printer/markdown_printer.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 import os
-from datetime import datetime
+from datetime import datetime, timezone
 
 from pyrit.executor.attack.printer.attack_result_printer import AttackResultPrinter
 from pyrit.memory import CentralMemory
@@ -171,7 +171,9 @@ async def print_result_async(
 
         # Footer
         markdown_lines.append("\n---")
-        markdown_lines.append(f"*Report generated at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*")
+        markdown_lines.append(
+            f"*Report generated at {datetime.now(tz=timezone.utc).strftime('%Y-%m-%d %H:%M:%S')} UTC*"
+        )
 
         self._render_markdown(markdown_lines)
 

pyrit/memory/memory_interface.py

@@ -9,7 +9,7 @@
 import weakref
 from collections.abc import MutableSequence, Sequence
 from contextlib import closing
-from datetime import datetime
+from datetime import datetime, timezone
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 
@@ -1001,7 +1001,7 @@ async def add_seeds_to_memory_async(self, *, seeds: Sequence[Seed], added_by: Op
             ValueError: If the 'added_by' attribute is not set for each prompt.
         """
         entries: MutableSequence[SeedEntry] = []
-        current_time = datetime.now()
+        current_time = datetime.now(tz=timezone.utc)
         for prompt in seeds:
             if added_by:
                 prompt.added_by = added_by
@@ -1245,7 +1245,7 @@ def export_conversations(
 
         # If file_path is not provided, construct a default using the exporter's results_path
         if not file_path:
-            file_name = f"exported_conversations_on_{datetime.now().strftime('%Y_%m_%d')}.{export_type}"
+            file_name = f"exported_conversations_on_{datetime.now(tz=timezone.utc).strftime('%Y_%m_%d')}.{export_type}"
             file_path = DB_DATA_PATH / file_name
 
         self.exporter.export_data(list(data), file_path=file_path, export_type=export_type)

pyrit/memory/memory_models.py

@@ -4,7 +4,7 @@
 import json
 import logging
 import uuid
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Literal, Optional, Union
 
 from pydantic import BaseModel, ConfigDict
@@ -59,6 +59,21 @@
 MAX_IDENTIFIER_VALUE_LENGTH: int = 80
 
 
+def _ensure_utc(dt: Optional[datetime]) -> Optional[datetime]:
+    """
+    Attach UTC tzinfo to a naive datetime (as returned by SQLite).
+
+    Args:
+        dt (Optional[datetime]): The datetime to normalize, or None.
+
+    Returns:
+        Optional[datetime]: The datetime with UTC tzinfo attached if it was naive, or None.
+    """
+    if dt is not None and dt.tzinfo is None:
+        return dt.replace(tzinfo=timezone.utc)
+    return dt
+
+
 class CustomUUID(TypeDecorator[uuid.UUID]):
     """
     A custom UUID type that works consistently across different database backends.
@@ -291,7 +306,7 @@ def get_message_piece(self) -> MessagePiece:
             converted_value_data_type=self.converted_value_data_type,
             response_error=self.response_error,
             original_prompt_id=self.original_prompt_id,
-            timestamp=self.timestamp,
+            timestamp=_ensure_utc(self.timestamp),
         )
         message_piece.scores = [score.get_score() for score in self.scores]
         return message_piece
@@ -416,7 +431,7 @@ def get_score(self) -> Score:
             score_metadata=self.score_metadata,
             scorer_class_identifier=scorer_identifier,
             message_piece_id=self.prompt_request_response_id,
-            timestamp=self.timestamp,
+            timestamp=_ensure_utc(self.timestamp),
             objective=self.objective,
         )
 
@@ -621,7 +636,7 @@ def get_seed(self) -> Seed:
                 authors=self.authors,
                 groups=self.groups,
                 source=self.source,
-                date_added=self.date_added,
+                date_added=_ensure_utc(self.date_added),
                 added_by=self.added_by,
                 metadata=self.prompt_metadata,
                 prompt_group_id=self.prompt_group_id,
@@ -641,7 +656,7 @@ def get_seed(self) -> Seed:
                 authors=self.authors,
                 groups=self.groups,
                 source=self.source,
-                date_added=self.date_added,
+                date_added=_ensure_utc(self.date_added),
                 added_by=self.added_by,
                 metadata=self.prompt_metadata,
                 prompt_group_id=self.prompt_group_id,
@@ -663,7 +678,7 @@ def get_seed(self) -> Seed:
             authors=self.authors,
             groups=self.groups,
             source=self.source,
-            date_added=self.date_added,
+            date_added=_ensure_utc(self.date_added),
             added_by=self.added_by,
             metadata=self.prompt_metadata,
             parameters=self.parameters,
@@ -774,7 +789,7 @@ def __init__(self, *, entry: AttackResult):
             ref.conversation_id for ref in entry.get_conversations_by_type(ConversationType.ADVERSARIAL)
         ] or None
 
-        self.timestamp = datetime.now()
+        self.timestamp = datetime.now(tz=timezone.utc)
         self.pyrit_version = pyrit.__version__
 
     @staticmethod
@@ -957,7 +972,7 @@ def __init__(self, *, entry: ScenarioResult):
             serialized_attack_results[attack_name] = [result.conversation_id for result in results]
         self.attack_results_json = json.dumps(serialized_attack_results)
 
-        self.timestamp = datetime.now()
+        self.timestamp = datetime.now(tz=timezone.utc)
 
     def get_scenario_result(self) -> ScenarioResult:
         """

pyrit/models/message.py

@@ -6,7 +6,7 @@
 import copy
 import uuid
 import warnings
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Optional, Union
 
 from pyrit.common.utils import combine_dict
@@ -415,7 +415,7 @@ def duplicate_message(self) -> Message:
 
         """
         new_pieces = copy.deepcopy(self.message_pieces)
-        new_timestamp = datetime.now()
+        new_timestamp = datetime.now(tz=timezone.utc)
         for piece in new_pieces:
             piece.id = uuid.uuid4()
             piece.timestamp = new_timestamp

pyrit/models/message_piece.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import uuid
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union, get_args
 from uuid import uuid4
 
@@ -108,7 +108,7 @@ def __init__(
         self.conversation_id = conversation_id if conversation_id else str(uuid4())
         self.sequence = sequence
 
-        self.timestamp = timestamp if timestamp else datetime.now()
+        self.timestamp = timestamp if timestamp else datetime.now(tz=timezone.utc)
         self.labels = labels or {}
         self.prompt_metadata = prompt_metadata or {}
 

pyrit/models/score.py

@@ -5,7 +5,7 @@
 
 import uuid
 from dataclasses import dataclass
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union, get_args
 
 if TYPE_CHECKING:
@@ -90,7 +90,7 @@ def __init__(
         from pyrit.identifiers.component_identifier import ComponentIdentifier
 
         self.id = id if id else uuid.uuid4()
-        self.timestamp = timestamp if timestamp else datetime.now()
+        self.timestamp = timestamp if timestamp else datetime.now(tz=timezone.utc)
 
         self.validate(score_type, score_value)
 

pyrit/models/seeds/seed.py

@@ -14,7 +14,7 @@
 import re
 import uuid
 from dataclasses import dataclass, field
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 
 from jinja2 import BaseLoader, Environment, StrictUndefined, Template, Undefined
@@ -113,7 +113,7 @@ class Seed(YamlLoadable):
     source: Optional[str] = None
 
     # Date when the prompt was added to the dataset
-    date_added: Optional[datetime] = field(default_factory=lambda: datetime.now())
+    date_added: Optional[datetime] = field(default_factory=lambda: datetime.now(tz=timezone.utc))
 
     # User who added the prompt to the dataset
     added_by: Optional[str] = None

pyrit/models/seeds/seed_dataset.py

@@ -12,7 +12,7 @@
 import uuid
 import warnings
 from collections import defaultdict
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 from pyrit.common import utils
@@ -121,7 +121,7 @@ def __init__(
         self.authors = authors or []
         self.groups = groups or []
         self.source = source
-        self.date_added = date_added or datetime.now()
+        self.date_added = date_added or datetime.now(tz=timezone.utc)
         self.added_by = added_by
 
         # Convert any dictionaries in `seeds` to SeedPrompt and/or SeedObjective objects