fix: include missing warn for Poetry (#758)

The previous warning implementation missed Poetry; this fix unifies the behavior.

Author: yeisonvargasf

safety/tool/mixins.py

@@ -0,0 +1,151 @@
+from typing import Any, List, Protocol, Tuple, Dict, Optional, runtime_checkable
+import typer
+from rich.padding import Padding
+
+from .base import EnvironmentDiffTracker
+
+from safety.console import main_console as console
+from safety.init.render import render_header, progressive_print
+from safety.models import ToolResult
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+@runtime_checkable
+class AuditableCommand(Protocol):
+    """
+    Protocol defining the contract for classes that can be audited for packages.
+    """
+
+    @property
+    def _diff_tracker(self) -> "EnvironmentDiffTracker":
+        """
+        Provides package tracking functionality.
+        """
+        ...
+
+    @property
+    def _packages(self) -> List[Tuple[str, Optional[str]]]:
+        """
+        Provides the target package list.
+        """
+        ...
+
+
+class InstallationAuditMixin:
+    """
+    Mixin providing installation audit functionality for command classes.
+
+    This mixin can be used by any command class that needs to audit
+    installation and show warnings.
+
+    Classes using this mixin should conform to the AuditableCommand protocol.
+    """
+
+    def render_installation_warnings(
+        self, ctx: typer.Context, packages_audit: Dict[str, Any]
+    ):
+        """
+        Render installation warnings based on package audit results.
+
+        Args:
+            ctx: The typer context
+            packages_audit: pre-fetched audit data
+        """
+
+        warning_messages = []
+        for audited_package in packages_audit.get("audit", {}).get("packages", []):
+            vulnerabilities = audited_package.get("vulnerabilities", {})
+            critical_vulnerabilities = vulnerabilities.get("critical", 0)
+            total_vulnerabilities = 0
+            for count in vulnerabilities.values():
+                total_vulnerabilities += count
+
+            if total_vulnerabilities == 0:
+                continue
+
+            warning_message = f"[[yellow]Warning[/yellow]] {audited_package.get('package_specifier')} contains {total_vulnerabilities} vulnerabilities"
+            if critical_vulnerabilities > 0:
+                warning_message += f", including {critical_vulnerabilities} critical severity vulnerabilities"
+
+            warning_message += "."
+            warning_messages.append(warning_message)
+
+        if len(warning_messages) > 0:
+            console.print()
+            render_header(" Safety Report")
+            progressive_print(warning_messages)
+            console.line()
+
+    def render_package_details(self: "AuditableCommand"):
+        """
+        Render details for installed packages.
+        """
+        for package_name, _ in self._packages:
+            console.print(
+                Padding(
+                    f"Learn more: [link]https://data.safetycli.com/packages/pypi/{package_name}/[/link]",
+                    (0, 0, 0, 1),
+                ),
+                emoji=True,
+            )
+
+    def audit_packages(self, ctx: typer.Context) -> Dict[str, Any]:
+        """
+        Audit packages based on environment diff tracking.
+        Override this method in your command class if needed.
+
+        Args:
+            ctx: The typer context
+
+        Returns:
+            Dict containing audit results
+        """
+        try:
+            # Check if the instance has a diff tracker and can get a diff
+            # Using getattr to avoid lint errors
+            diff_tracker = getattr(self, "_diff_tracker", None)
+            if diff_tracker and hasattr(diff_tracker, "get_diff"):
+                added, _, updated = diff_tracker.get_diff()
+                packages = {**added, **updated}
+
+                if hasattr(ctx.obj, "auth") and hasattr(ctx.obj.auth, "client"):
+                    return ctx.obj.auth.client.audit_packages(
+                        [
+                            f"{package_name}=={version[-1] if isinstance(version, tuple) else version}"
+                            for (package_name, version) in packages.items()
+                        ]
+                    )
+        except Exception:
+            logger.debug("Audit API failed with error", exc_info=True)
+
+        # Always return a dict to satisfy the return type
+        return dict()
+
+    def handle_installation_audit(self, ctx: typer.Context, result: ToolResult):
+        """
+        Handle installation audit and rendering warnings/details.
+        This is an explicit method that can be called from a command's after method.
+
+        Usage example:
+            def after(self, ctx, result):
+                super().after(ctx, result)
+                self.handle_installation_audit(ctx, result)
+
+        Args:
+            ctx: The typer context
+            result: The tool result
+        """
+
+        if not isinstance(self, AuditableCommand):
+            raise TypeError(
+                "handle_installation_audit can only be called on instances of AuditableCommand"
+            )
+
+        packages_audit = self.audit_packages(ctx)
+        self.render_installation_warnings(ctx, packages_audit)
+
+        # If command failed, show package details
+        if not result.process or result.process.returncode != 0:
+            self.render_package_details()

safety/tool/pip/command.py

@@ -2,10 +2,9 @@
 from pathlib import Path
 import re
 from tempfile import mkstemp
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, List, Optional
 
 import logging
-from rich.padding import Padding
 import typer
 
 from safety.models import ToolResult
@@ -15,10 +14,9 @@
 from ..intents import ToolIntentionType
 from safety_schemas.models.events.types import ToolType
 from ..environment_diff import EnvironmentDiffTracker, PipEnvironmentDiffTracker
+from ..mixins import InstallationAuditMixin
 from ..utils import Pip
 
-from safety.console import main_console as console
-from ...init.render import render_header, progressive_print
 
 PIP_LOCK = "safety-pip.lock"
 
@@ -80,10 +78,10 @@ class PipGenericCommand(PipCommand):
     pass
 
 
-class PipInstallCommand(PipCommand):
+class PipInstallCommand(PipCommand, InstallationAuditMixin):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        self.__packages = []
+        self._packages = []
         self.__index_url = None
 
     def before(self, ctx: typer.Context):
@@ -92,7 +90,7 @@ def before(self, ctx: typer.Context):
 
         if self._intention:
             for pkg in self._intention.packages:
-                self.__packages.append((pkg.name, pkg.version_constraint))
+                self._packages.append((pkg.name, pkg.version_constraint))
 
             if index_opt := self._intention.options.get(
                 "index-url"
@@ -137,65 +135,9 @@ def before(self, ctx: typer.Context):
 
     def after(self, ctx: typer.Context, result: ToolResult):
         super().after(ctx, result)
-
-        self.__render_installation_warnings(ctx)
-
-        if not result.process or result.process.returncode != 0:
-            self.__render_package_details()
+        self.handle_installation_audit(ctx, result)
 
     def env(self, ctx: typer.Context) -> dict:
         env = super().env(ctx)
         env["PIP_INDEX_URL"] = Pip.build_index_url(ctx, self.__index_url)
         return env
-
-    def __render_installation_warnings(self, ctx: typer.Context):
-        packages_audit = self.__audit_packages(ctx)
-
-        warning_messages = []
-        for audited_package in packages_audit.get("audit", {}).get("packages", []):
-            vulnerabilities = audited_package.get("vulnerabilities", {})
-            critical_vulnerabilities = vulnerabilities.get("critical", 0)
-            total_vulnerabilities = 0
-            for count in vulnerabilities.values():
-                total_vulnerabilities += count
-
-            if total_vulnerabilities == 0:
-                continue
-
-            warning_message = f"[[yellow]Warning[/yellow]] {audited_package.get('package_specifier')} contains {total_vulnerabilities} vulnerabilities"
-            if critical_vulnerabilities > 0:
-                warning_message += f", including {critical_vulnerabilities} critical severity vulnerabilities"
-
-            warning_message += "."
-            warning_messages.append(warning_message)
-
-        if len(warning_messages) > 0:
-            console.print()
-            render_header(" Safety Report")
-            progressive_print(warning_messages)
-
-    def __render_package_details(self):
-        for package_name, version_specifier in self.__packages:
-            console.print(
-                Padding(
-                    f"Learn more: [link]https://data.safetycli.com/packages/pypi/{package_name}/[/link]",
-                    (0, 0, 0, 1),
-                ),
-                emoji=True,
-            )
-
-    def __audit_packages(self, ctx: typer.Context) -> Any:
-        try:
-            added, _, updated = self._diff_tracker.get_diff()
-            packages = {**added, **updated}
-
-            return ctx.obj.auth.client.audit_packages(
-                [
-                    f"{package_name}=={version[-1] if isinstance(version, tuple) else version}"
-                    for (package_name, version) in packages.items()
-                ]
-            )
-        except Exception:
-            logger.debug("Audit API failed with error", exc_info=True)
-            # do not propagate the error in case the audit failed
-            return dict()

safety/tool/poetry/command.py

@@ -11,10 +11,12 @@
 
 from ..auth import index_credentials
 from ..base import BaseCommand, ToolIntentionType
+from ..mixins import InstallationAuditMixin
 from ..environment_diff import EnvironmentDiffTracker, PipEnvironmentDiffTracker
 from safety_schemas.models.events.types import ToolType
 
 from safety.console import main_console as console
+from safety.models import ToolResult
 
 PO_LOCK = "safety-po.lock"
 
@@ -161,7 +163,11 @@ class PoetryGenericCommand(PoetryCommand):
     pass
 
 
-class PoetryAddCommand(PoetryCommand):
+class PoetryAddCommand(PoetryCommand, InstallationAuditMixin):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._packages = []
+
     def patch_source_option(
         self, args: List[str], new_source: str = "safety"
     ) -> Tuple[Optional[str], List[str]]:
@@ -197,3 +203,19 @@ def before(self, ctx: typer.Context):
 
         _, modified_args = self.patch_source_option(self._args)
         self._args = modified_args
+
+        # Extract packages from intention for rendering later
+        if self._intention and self._intention.packages:
+            for pkg in self._intention.packages:
+                self._packages.append((pkg.name, pkg.version_constraint))
+
+    def after(self, ctx: typer.Context, result: ToolResult):
+        """
+        Run after the command execution. Handle installation audit via mixin.
+
+        Args:
+            ctx: The typer context
+            result: The tool result
+        """
+        super().after(ctx, result)
+        self.handle_installation_audit(ctx, result)

tests/tool/test_installation_commands.py

@@ -0,0 +1,48 @@
+# type: ignore
+
+import pytest
+from unittest.mock import MagicMock, patch
+
+import typer
+
+from safety.tool.pip.command import PipInstallCommand
+from safety.tool.uv.command import UvInstallCommand
+from safety.tool.poetry.command import PoetryAddCommand
+
+
+class TestInstallationCommandsAudit:
+    """
+    Test suite for verifying installation audit functionality in command classes.
+    """
+
+    def setup_method(self):
+        """
+        Set up test fixtures.
+        """
+        self.ctx = MagicMock(spec=typer.Context)
+        self.ctx.obj = MagicMock()
+        self.result = MagicMock(duration_ms=100, process=MagicMock(returncode=0))
+
+    @pytest.mark.parametrize(
+        "command_class,command_args",
+        [
+            (PipInstallCommand, ["install", "requests"]),
+            (UvInstallCommand, ["pip", "install", "requests"]),
+            (PoetryAddCommand, ["add", "requests"]),
+        ],
+    )
+    @patch("safety.tool.base.BaseCommand._handle_command_result")
+    def test_installation_command_calls_audit(
+        self, mock_handle_result, command_class, command_args
+    ):
+        """
+        Test that all installation commands call handle_installation_audit in after().
+        """
+        command = command_class(command_args)
+
+        with patch.object(
+            command_class, "handle_installation_audit"
+        ) as mock_handle_audit:
+            command.after(self.ctx, self.result)
+
+            mock_handle_audit.assert_called_once_with(self.ctx, self.result)