From 9a135262c558507e32e952e2b5484bb4eea1c6b3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 18 Apr 2025 03:31:27 +0000 Subject: [PATCH] Fix lint and type-checker issues Co-Authored-By: Joe Moura --- src/crewai/cli/plot_flow.py | 19 +++++++++++-------- tests/test_crew_plot.py | 2 +- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/crewai/cli/plot_flow.py b/src/crewai/cli/plot_flow.py index b8afae4c6..3e61d8694 100644 --- a/src/crewai/cli/plot_flow.py +++ b/src/crewai/cli/plot_flow.py @@ -1,5 +1,5 @@ -import os import importlib.util +import os import sys import click @@ -14,14 +14,17 @@ def plot_flow() -> None: try: if os.path.exists(main_file_path): spec = importlib.util.spec_from_file_location("main", main_file_path) - main = importlib.util.module_from_spec(spec) - sys.modules["main"] = main - spec.loader.exec_module(main) - - if hasattr(main, "plot"): - main.plot() + if spec is not None and spec.loader is not None: + main = importlib.util.module_from_spec(spec) + sys.modules["main"] = main + spec.loader.exec_module(main) + + if hasattr(main, "plot"): + main.plot() + else: + click.echo("Error: No plot function found in main.py", err=True) else: - click.echo("Error: No plot function found in main.py", err=True) + click.echo("Error: Could not load main.py module", err=True) else: click.echo("Error: Could not find main.py in the current directory", err=True) diff --git a/tests/test_crew_plot.py b/tests/test_crew_plot.py index dc833a015..3fc1b1223 100644 --- a/tests/test_crew_plot.py +++ b/tests/test_crew_plot.py @@ -1,7 +1,7 @@ import pytest -from crewai.crew import Crew from crewai.agent import Agent +from crewai.crew import Crew def test_crew_plot_method_not_implemented():