chore: fix ruff linting and mypy issues in flow module

Author: greysonlalonde

src/crewai/flow/__init__.py

@@ -1,5 +1,4 @@
-from crewai.flow.flow import Flow, start, listen, or_, and_, router
+from crewai.flow.flow import Flow, and_, listen, or_, router, start
 from crewai.flow.persistence import persist
 
-__all__ = ["Flow", "start", "listen", "or_", "and_", "router", "persist"]
-
+__all__ = ["Flow", "and_", "listen", "or_", "persist", "router", "start"]

src/crewai/flow/flow_trackable.py

@@ -1,5 +1,4 @@
 import inspect
-from typing import Optional
 
 from pydantic import BaseModel, Field, InstanceOf, model_validator
 
@@ -14,7 +13,7 @@ class FlowTrackable(BaseModel):
     inspecting the call stack.
     """
 
-    parent_flow: Optional[InstanceOf[Flow]] = Field(
+    parent_flow: InstanceOf[Flow] | None = Field(
         default=None,
         description="The parent flow of the instance, if it was created inside a flow.",
     )

src/crewai/flow/flow_visualizer.py

@@ -1,14 +1,13 @@
 # flow_visualizer.py
 
 import os
-from pathlib import Path
 
-from pyvis.network import Network
+from pyvis.network import Network  # type: ignore[import-untyped]
 
 from crewai.flow.config import COLORS, NODE_STYLES
 from crewai.flow.html_template_handler import HTMLTemplateHandler
 from crewai.flow.legend_generator import generate_legend_items_html, get_legend_items
-from crewai.flow.path_utils import safe_path_join, validate_path_exists
+from crewai.flow.path_utils import safe_path_join
 from crewai.flow.utils import calculate_node_levels
 from crewai.flow.visualization_utils import (
     add_edges,
@@ -34,13 +33,13 @@ def __init__(self, flow):
         ValueError
             If flow object is invalid or missing required attributes.
         """
-        if not hasattr(flow, '_methods'):
+        if not hasattr(flow, "_methods"):
             raise ValueError("Invalid flow object: missing '_methods' attribute")
-        if not hasattr(flow, '_listeners'):
+        if not hasattr(flow, "_listeners"):
             raise ValueError("Invalid flow object: missing '_listeners' attribute")
-        if not hasattr(flow, '_start_methods'):
+        if not hasattr(flow, "_start_methods"):
             raise ValueError("Invalid flow object: missing '_start_methods' attribute")
-            
+
         self.flow = flow
         self.colors = COLORS
         self.node_styles = NODE_STYLES
@@ -65,7 +64,7 @@ def plot(self, filename):
         """
         if not filename or not isinstance(filename, str):
             raise ValueError("Filename must be a non-empty string")
-            
+
         try:
             # Initialize network
             net = Network(
@@ -96,45 +95,51 @@ def plot(self, filename):
             try:
                 node_levels = calculate_node_levels(self.flow)
             except Exception as e:
-                raise ValueError(f"Failed to calculate node levels: {str(e)}")
+                raise ValueError(f"Failed to calculate node levels: {e!s}") from e
 
             # Compute positions
             try:
                 node_positions = compute_positions(self.flow, node_levels)
             except Exception as e:
-                raise ValueError(f"Failed to compute node positions: {str(e)}")
+                raise ValueError(f"Failed to compute node positions: {e!s}") from e
 
             # Add nodes to the network
             try:
                 add_nodes_to_network(net, self.flow, node_positions, self.node_styles)
             except Exception as e:
-                raise RuntimeError(f"Failed to add nodes to network: {str(e)}")
+                raise RuntimeError(f"Failed to add nodes to network: {e!s}") from e
 
             # Add edges to the network
             try:
                 add_edges(net, self.flow, node_positions, self.colors)
             except Exception as e:
-                raise RuntimeError(f"Failed to add edges to network: {str(e)}")
+                raise RuntimeError(f"Failed to add edges to network: {e!s}") from e
 
             # Generate HTML
             try:
                 network_html = net.generate_html()
                 final_html_content = self._generate_final_html(network_html)
             except Exception as e:
-                raise RuntimeError(f"Failed to generate network visualization: {str(e)}")
+                raise RuntimeError(
+                    f"Failed to generate network visualization: {e!s}"
+                ) from e
 
             # Save the final HTML content to the file
             try:
                 with open(f"{filename}.html", "w", encoding="utf-8") as f:
                     f.write(final_html_content)
                 print(f"Plot saved as {filename}.html")
             except IOError as e:
-                raise IOError(f"Failed to save flow visualization to {filename}.html: {str(e)}")
+                raise IOError(
+                    f"Failed to save flow visualization to {filename}.html: {e!s}"
+                ) from e
 
         except (ValueError, RuntimeError, IOError) as e:
             raise e
         except Exception as e:
-            raise RuntimeError(f"Unexpected error during flow visualization: {str(e)}")
+            raise RuntimeError(
+                f"Unexpected error during flow visualization: {e!s}"
+            ) from e
         finally:
             self._cleanup_pyvis_lib()
 
@@ -165,7 +170,9 @@ def _generate_final_html(self, network_html):
         try:
             # Extract just the body content from the generated HTML
             current_dir = os.path.dirname(__file__)
-            template_path = safe_path_join("assets", "crewai_flow_visual_template.html", root=current_dir)
+            template_path = safe_path_join(
+                "assets", "crewai_flow_visual_template.html", root=current_dir
+            )
             logo_path = safe_path_join("assets", "crewai_logo.svg", root=current_dir)
 
             if not os.path.exists(template_path):
@@ -179,12 +186,9 @@ def _generate_final_html(self, network_html):
             # Generate the legend items HTML
             legend_items = get_legend_items(self.colors)
             legend_items_html = generate_legend_items_html(legend_items)
-            final_html_content = html_handler.generate_final_html(
-                network_body, legend_items_html
-            )
-            return final_html_content
+            return html_handler.generate_final_html(network_body, legend_items_html)
         except Exception as e:
-            raise IOError(f"Failed to generate visualization HTML: {str(e)}")
+            raise IOError(f"Failed to generate visualization HTML: {e!s}") from e
 
     def _cleanup_pyvis_lib(self):
         """
@@ -197,6 +201,7 @@ def _cleanup_pyvis_lib(self):
             lib_folder = safe_path_join("lib", root=os.getcwd())
             if os.path.exists(lib_folder) and os.path.isdir(lib_folder):
                 import shutil
+
                 shutil.rmtree(lib_folder)
         except ValueError as e:
             print(f"Error validating lib folder path: {e}")

src/crewai/flow/html_template_handler.py

@@ -1,8 +1,7 @@
 import base64
 import re
-from pathlib import Path
 
-from crewai.flow.path_utils import safe_path_join, validate_path_exists
+from crewai.flow.path_utils import validate_path_exists
 
 
 class HTMLTemplateHandler:
@@ -28,7 +27,7 @@ def __init__(self, template_path, logo_path):
             self.template_path = validate_path_exists(template_path, "file")
             self.logo_path = validate_path_exists(logo_path, "file")
         except ValueError as e:
-            raise ValueError(f"Invalid template or logo path: {e}")
+            raise ValueError(f"Invalid template or logo path: {e}") from e
 
     def read_template(self):
         """Read and return the HTML template file contents."""
@@ -53,23 +52,23 @@ def generate_legend_items_html(self, legend_items):
             if "border" in item:
                 legend_items_html += f"""
                 <div class="legend-item">
-                <div class="legend-color-box" style="background-color: {item['color']}; border: 2px dashed {item['border']};"></div>
-                <div>{item['label']}</div>
+                <div class="legend-color-box" style="background-color: {item["color"]}; border: 2px dashed {item["border"]};"></div>
+                <div>{item["label"]}</div>
                 </div>
                 """
             elif item.get("dashed") is not None:
                 style = "dashed" if item["dashed"] else "solid"
                 legend_items_html += f"""
                 <div class="legend-item">
-                <div class="legend-{style}" style="border-bottom: 2px {style} {item['color']};"></div>
-                <div>{item['label']}</div>
+                <div class="legend-{style}" style="border-bottom: 2px {style} {item["color"]};"></div>
+                <div>{item["label"]}</div>
                 </div>
                 """
             else:
                 legend_items_html += f"""
                 <div class="legend-item">
-                <div class="legend-color-box" style="background-color: {item['color']};"></div>
-                <div>{item['label']}</div>
+                <div class="legend-color-box" style="background-color: {item["color"]};"></div>
+                <div>{item["label"]}</div>
                 </div>
                 """
         return legend_items_html
@@ -79,15 +78,9 @@ def generate_final_html(self, network_body, legend_items_html, title="Flow Plot"
         html_template = self.read_template()
         logo_svg_base64 = self.encode_logo()
 
-        final_html_content = html_template.replace("{{ title }}", title)
-        final_html_content = final_html_content.replace(
-            "{{ network_content }}", network_body
+        return (
+            html_template.replace("{{ title }}", title)
+            .replace("{{ network_content }}", network_body)
+            .replace("{{ logo_svg_base64 }}", logo_svg_base64)
+            .replace("<!-- LEGEND_ITEMS_PLACEHOLDER -->", legend_items_html)
         )
-        final_html_content = final_html_content.replace(
-            "{{ logo_svg_base64 }}", logo_svg_base64
-        )
-        final_html_content = final_html_content.replace(
-            "<!-- LEGEND_ITEMS_PLACEHOLDER -->", legend_items_html
-        )
-
-        return final_html_content

src/crewai/flow/path_utils.py

@@ -5,12 +5,10 @@
 traversal attacks and ensure paths remain within allowed boundaries.
 """
 
-import os
 from pathlib import Path
-from typing import List, Union
 
 
-def safe_path_join(*parts: str, root: Union[str, Path, None] = None) -> str:
+def safe_path_join(*parts: str, root: str | Path | None = None) -> str:
     """
     Safely join path components and ensure the result is within allowed boundaries.
 
@@ -43,25 +41,25 @@ def safe_path_join(*parts: str, root: Union[str, Path, None] = None) -> str:
 
         # Establish root directory
         root_path = Path(root).resolve() if root else Path.cwd()
-        
+
         # Join and resolve the full path
         full_path = Path(root_path, *clean_parts).resolve()
-        
+
         # Check if the resolved path is within root
         if not str(full_path).startswith(str(root_path)):
             raise ValueError(
                 f"Invalid path: Potential directory traversal. Path must be within {root_path}"
             )
-            
+
         return str(full_path)
-        
+
     except Exception as e:
         if isinstance(e, ValueError):
             raise
-        raise ValueError(f"Invalid path components: {str(e)}")
+        raise ValueError(f"Invalid path components: {e!s}") from e
 
 
-def validate_path_exists(path: Union[str, Path], file_type: str = "file") -> str:
+def validate_path_exists(path: str | Path, file_type: str = "file") -> str:
     """
     Validate that a path exists and is of the expected type.
 
@@ -84,24 +82,24 @@ def validate_path_exists(path: Union[str, Path], file_type: str = "file") -> str
     """
     try:
         path_obj = Path(path).resolve()
-        
+
         if not path_obj.exists():
             raise ValueError(f"Path does not exist: {path}")
-            
+
         if file_type == "file" and not path_obj.is_file():
             raise ValueError(f"Path is not a file: {path}")
-        elif file_type == "directory" and not path_obj.is_dir():
+        if file_type == "directory" and not path_obj.is_dir():
             raise ValueError(f"Path is not a directory: {path}")
-            
+
         return str(path_obj)
-        
+
     except Exception as e:
         if isinstance(e, ValueError):
             raise
-        raise ValueError(f"Invalid path: {str(e)}")
+        raise ValueError(f"Invalid path: {e!s}") from e
 
 
-def list_files(directory: Union[str, Path], pattern: str = "*") -> List[str]:
+def list_files(directory: str | Path, pattern: str = "*") -> list[str]:
     """
     Safely list files in a directory matching a pattern.
 
@@ -126,10 +124,10 @@ def list_files(directory: Union[str, Path], pattern: str = "*") -> List[str]:
         dir_path = Path(directory).resolve()
         if not dir_path.is_dir():
             raise ValueError(f"Not a directory: {directory}")
-            
+
         return [str(p) for p in dir_path.glob(pattern) if p.is_file()]
-        
+
     except Exception as e:
         if isinstance(e, ValueError):
             raise
-        raise ValueError(f"Error listing files: {str(e)}")
+        raise ValueError(f"Error listing files: {e!s}") from e

src/crewai/flow/persistence/__init__.py

@@ -4,15 +4,15 @@
 This module provides interfaces and implementations for persisting flow states.
 """
 
-from typing import Any, Dict, TypeVar, Union
+from typing import Any, TypeVar
 
 from pydantic import BaseModel
 
 from crewai.flow.persistence.base import FlowPersistence
 from crewai.flow.persistence.decorators import persist
 from crewai.flow.persistence.sqlite import SQLiteFlowPersistence
 
-__all__ = ["FlowPersistence", "persist", "SQLiteFlowPersistence"]
+__all__ = ["FlowPersistence", "SQLiteFlowPersistence", "persist"]
 
-StateType = TypeVar('StateType', bound=Union[Dict[str, Any], BaseModel])
-DictStateType = Dict[str, Any]
+StateType = TypeVar("StateType", bound=dict[str, Any] | BaseModel)
+DictStateType = dict[str, Any]