Enable DOC201 check and fix existing violations (#203)

* Add docstring check for DOC201

Added check to pyproject.toml, and fixed 46 docstrings where DOC201 was flagged.

Implements: #159

* Update pyproject to pass pre-commit.

Previous change for #159 added `preview = true` to pyproject.toml, which added unstable preview rules.
Fixed this by adding `explicit-preview-rules = true`.

* Update pyproject.toml to move explicit-preview-rules to ruff.lint.

---------

Co-authored-by: Hamza Sardar <[email protected]>

Author: HamzaSardar

pyproject.toml

@@ -79,6 +79,7 @@ indent-width = 4
 
 target-version = "py312"
 force-exclude = true
+preview = true
 
 [tool.ruff.lint]
 select = [
@@ -96,15 +97,17 @@ select = [
     "PLE", # ruff currently implements only a subset of pylint's rules
     "PLW", # pylint warning
     "PLR", # pylint refactor
-    "UP",  # pyupgrade
-    "C",   # Complexity (mccabe+) & comprehensions
+    "UP",    # pyupgrade
+    "C",     # Complexity (mccabe+) & comprehensions
+    "DOC201", # pydoclint: return value must be documented
 ]
 ignore = [
     "UP006",   # See https://github.com/bokeh/bokeh/issues/13143
     "UP007",   # See https://github.com/bokeh/bokeh/pull/13144
     "PLC0415", # Allow imports inside functions (useful for optional deps)
     "PLR2004", # Allow magic values in comparisons (array indices etc.)
 ]
+explicit-preview-rules = true
 
 [tool.ruff.format]
 # Like Black, use double quotes for strings.

src/sre_agent/cli/configuration/wizard.py

@@ -189,7 +189,11 @@ def _configure_model_provider(
     updates: dict[str, str],
     allow_back: bool = False,
 ) -> str:
-    """Prompt for model provider and required credentials."""
+    """Prompt for model provider and required credentials.
+
+    Returns:
+        The selected model provider name.
+    """
     model_provider = _prompt_choice(
         "Model provider:",
         config.integrations.model_provider,
@@ -215,7 +219,11 @@ def _configure_notification_platform(
     updates: dict[str, str],
     allow_back: bool = False,
 ) -> tuple[str, str | None]:
-    """Prompt for notification platform and required credentials."""
+    """Prompt for notification platform and required credentials.
+
+    Returns:
+        A tuple of (notification_platform, slack_channel_id).
+    """
     notification_platform = _prompt_choice(
         "Messaging/notification platform:",
         config.integrations.notification_platform,
@@ -248,7 +256,11 @@ def _configure_code_repository_provider(
     updates: dict[str, str],
     allow_back: bool = False,
 ) -> tuple[str, str | None, str | None, str | None]:
-    """Prompt for code repository provider and required credentials."""
+    """Prompt for code repository provider and required credentials.
+
+    Returns:
+        A tuple of (code_repository_provider, github_owner, github_repo, github_ref).
+    """
     code_repository_provider = _prompt_choice(
         "Remote code repository:",
         config.integrations.code_repository_provider,
@@ -299,7 +311,11 @@ def _configure_deployment_platform(
     updates: dict[str, str],
     allow_back: bool = False,
 ) -> tuple[str, str]:
-    """Prompt for deployment platform, logging platform, and AWS credentials."""
+    """Prompt for deployment platform, logging platform, and AWS credentials.
+
+    Returns:
+        A tuple of (deployment_platform, logging_platform).
+    """
     deployment_platform = _prompt_choice(
         "Which platform is your application deployed on?",
         config.integrations.deployment_platform,

src/sre_agent/cli/presentation/banner.py

@@ -39,7 +39,11 @@ def _print_animated_banner() -> None:
 
 
 def _build_banner(colour_offset: int) -> Panel:
-    """Build the banner panel with a shifted colour palette."""
+    """Build the banner panel with a shifted colour palette.
+
+    Returns:
+        A Rich Panel containing the styled ASCII art banner.
+    """
     ascii_art = get_ascii_art().strip("\n")
     # spellchecker:ignore-next-line
     banner_text = Text(justify="center")

src/sre_agent/cli/presentation/styles.py

@@ -4,22 +4,20 @@
 import questionary.constants as questionary_constants
 import questionary.styles as questionary_styles
 
-QUESTIONARY_STYLE = questionary.Style(
-    [
-        ("qmark", "fg:#7C3AED"),
-        ("question", "fg:#e0e0e0 bold"),
-        ("answer", "fg:#5EEAD4 bold"),
-        ("search_success", "noinherit fg:#00FF00 bold"),
-        ("search_none", "noinherit fg:#FF0000 bold"),
-        ("pointer", "fg:#e0e0e0"),
-        ("highlighted", "fg:#f2f2f2"),
-        ("selected", "fg:#e0e0e0"),
-        ("separator", "fg:#e0e0e0"),
-        ("instruction", "fg:#e0e0e0"),
-        ("text", "fg:#e0e0e0"),
-        ("disabled", "fg:#bdbdbd italic"),
-    ]
-)
+QUESTIONARY_STYLE = questionary.Style([
+    ("qmark", "fg:#7C3AED"),
+    ("question", "fg:#e0e0e0 bold"),
+    ("answer", "fg:#5EEAD4 bold"),
+    ("search_success", "noinherit fg:#00FF00 bold"),
+    ("search_none", "noinherit fg:#FF0000 bold"),
+    ("pointer", "fg:#e0e0e0"),
+    ("highlighted", "fg:#f2f2f2"),
+    ("selected", "fg:#e0e0e0"),
+    ("separator", "fg:#e0e0e0"),
+    ("instruction", "fg:#e0e0e0"),
+    ("text", "fg:#e0e0e0"),
+    ("disabled", "fg:#bdbdbd italic"),
+])
 
 
 def apply_questionary_style() -> None:

src/sre_agent/core/deployments/aws_ecs/cleanup.py

@@ -278,7 +278,11 @@ def _wait_for_nat_gateways(ec2: Any, nat_ids: list[str], reporter: Callable[[str
 
 
 def _list_internet_gateways(ec2: Any, vpc_id: str) -> list[str]:
-    """List internet gateways attached to a VPC."""
+    """List internet gateways attached to a VPC.
+
+    Returns:
+        A list of internet gateway IDs attached to the VPC.
+    """
     response = ec2.describe_internet_gateways(
         Filters=[{"Name": "attachment.vpc-id", "Values": [vpc_id]}]
     )

src/sre_agent/core/deployments/aws_ecs/ecr.py

@@ -7,7 +7,11 @@
 
 
 def ensure_repository(session: Session, name: str) -> str:
-    """Ensure an ECR repository exists and return its URI."""
+    """Ensure an ECR repository exists and return its URI.
+
+    Returns:
+        The URI of the ECR repository.
+    """
     ecr = session.client("ecr")
     try:
         response = ecr.describe_repositories(repositoryNames=[name])

src/sre_agent/core/deployments/aws_ecs/ecs_tasks.py

@@ -29,7 +29,11 @@ def register_task_definition(
     config: EcsDeploymentConfig,
     reporter: Callable[[str], None],
 ) -> str:
-    """Register the ECS task definition."""
+    """Register the ECS task definition.
+
+    Returns:
+        The ARN of the registered task definition.
+    """
     cpu_architecture = _normalise_cpu_architecture(config.task_cpu_architecture)
     if not config.exec_role_arn or not config.task_role_arn:
         raise RuntimeError("Task roles must be created before registering the task definition.")
@@ -134,7 +138,11 @@ def _normalise_cpu_architecture(value: str) -> str:
 
 
 def ensure_cluster(session: Session, cluster_name: str) -> str:
-    """Ensure an ECS cluster exists."""
+    """Ensure an ECS cluster exists.
+
+    Returns:
+        The ARN of the ECS cluster.
+    """
     ecs = session.client("ecs")
     response = ecs.describe_clusters(clusters=[cluster_name])
     clusters = response.get("clusters", [])
@@ -159,7 +167,11 @@ def run_task(
     config: EcsDeploymentConfig,
     container_overrides: list[dict[str, Any]] | None = None,
 ) -> str:
-    """Run a one-off ECS task."""
+    """Run a one-off ECS task.
+
+    Returns:
+        The ARN of the launched ECS task.
+    """
     if not config.task_definition_arn:
         raise RuntimeError("Task definition is missing. Register it before running tasks.")
     if not config.security_group_id or not config.private_subnet_ids:
@@ -209,7 +221,11 @@ def wait_for_task_completion(
     timeout_seconds: int = 1800,
     poll_interval_seconds: int = 5,
 ) -> tuple[bool, str]:
-    """Wait for a task to stop and report container exit status."""
+    """Wait for a task to stop and report container exit status.
+
+    Returns:
+        A tuple of (success, message) indicating the task outcome.
+    """
     ecs = session.client("ecs")
     deadline = time.time() + timeout_seconds
 
@@ -232,7 +248,11 @@ def wait_for_task_completion(
 
 
 def _task_completion_result(task: dict[str, Any]) -> tuple[bool, str]:
-    """Convert ECS task details into a completion result."""
+    """Convert ECS task details into a completion result.
+
+    Returns:
+        A tuple of (success, message) indicating the task outcome.
+    """
     target = _find_container(task.get("containers", []), SRE_AGENT_CONTAINER_NAME)
     if target is None:
         stopped_reason = str(task.get("stoppedReason", "task stopped"))

src/sre_agent/core/deployments/aws_ecs/iam.py

@@ -15,7 +15,11 @@ def ensure_roles(
     secret_arns: list[str],
     reporter: Callable[[str], None],
 ) -> tuple[str, str]:
-    """Ensure execution and task roles exist."""
+    """Ensure execution and task roles exist.
+
+    Returns:
+        A tuple of (exec_role_arn, task_role_arn).
+    """
     if not secret_arns:
         raise RuntimeError("Secret ARNs are required before creating roles.")
 
@@ -71,7 +75,11 @@ def ensure_service_linked_role(session: Session, reporter: Callable[[str], None]
 
 
 def _ensure_role(iam: Any, role_name: str, trust_policy: dict[str, Any]) -> str:
-    """Create a role if needed and return its ARN."""
+    """Create a role if needed and return its ARN.
+
+    Returns:
+        The ARN of the IAM role.
+    """
     try:
         response = iam.get_role(RoleName=role_name)
         return cast(str, response["Role"]["Arn"])
@@ -125,7 +133,11 @@ def _ecs_trust_policy() -> dict[str, Any]:
 
 
 def _secrets_policy(secret_arns: list[str]) -> dict[str, Any]:
-    """Allow read access to Secrets Manager."""
+    """Allow read access to Secrets Manager.
+
+    Returns:
+        An IAM policy document granting read access to the given secrets.
+    """
     return {
         "Version": "2012-10-17",
         "Statement": [
@@ -139,7 +151,11 @@ def _secrets_policy(secret_arns: list[str]) -> dict[str, Any]:
 
 
 def _logs_policy(region: str, account_id: str) -> dict[str, Any]:
-    """Allow CloudWatch Logs queries."""
+    """Allow CloudWatch Logs queries.
+
+    Returns:
+        An IAM policy document granting CloudWatch Logs query access.
+    """
     return {
         "Version": "2012-10-17",
         "Statement": [

src/sre_agent/core/deployments/aws_ecs/images.py

@@ -28,7 +28,11 @@ def build_and_push_images(
     image_config: ImageBuildConfig,
     reporter: Callable[[str], None],
 ) -> str:
-    """Build and push container images to ECR."""
+    """Build and push container images to ECR.
+
+    Returns:
+        The ECS runtime CPU architecture used for the build.
+    """
     _require_docker()
 
     reporter("Authenticating Docker with ECR")

src/sre_agent/core/deployments/aws_ecs/network.py

@@ -13,7 +13,11 @@ def create_basic_vpc(
     project_name: str,
     reporter: Callable[[str], None],
 ) -> NetworkSelection:
-    """Create a simple VPC with one public and one private subnet."""
+    """Create a simple VPC with one public and one private subnet.
+
+    Returns:
+        A NetworkSelection containing the VPC ID and private subnet IDs.
+    """
     ec2 = session.client("ec2")
 
     reporter("Creating VPC (private networking foundation)")
@@ -91,7 +95,11 @@ def _tag_resource(ec2: Any, resource_id: str, name: str) -> None:
 
 
 def _first_availability_zone(ec2: Any) -> str:
-    """Fetch the first availability zone."""
+    """Fetch the first availability zone.
+
+    Returns:
+        The name of the first available availability zone.
+    """
     response = ec2.describe_availability_zones()
     zones = response.get("AvailabilityZones", [])
     if not zones: