mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
feat: mypy plugin base
* feat: base mypy plugin with CrewBase * fix: add crew method to protocol
This commit is contained in:
60
lib/crewai/src/crewai/mypy.py
Normal file
60
lib/crewai/src/crewai/mypy.py
Normal file
@@ -0,0 +1,60 @@
|
|||||||
|
"""Mypy plugin for CrewAI decorator type checking.
|
||||||
|
|
||||||
|
This plugin informs mypy about attributes injected by the @CrewBase decorator.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from collections.abc import Callable
|
||||||
|
|
||||||
|
from mypy.nodes import MDEF, SymbolTableNode, Var
|
||||||
|
from mypy.plugin import ClassDefContext, Plugin
|
||||||
|
from mypy.types import AnyType, TypeOfAny
|
||||||
|
|
||||||
|
|
||||||
|
class CrewAIPlugin(Plugin):
|
||||||
|
"""Mypy plugin that handles @CrewBase decorator attribute injection."""
|
||||||
|
|
||||||
|
def get_class_decorator_hook(
|
||||||
|
self, fullname: str
|
||||||
|
) -> Callable[[ClassDefContext], None] | None:
|
||||||
|
"""Return hook for class decorators.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
fullname: Fully qualified name of the decorator.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Hook function if this is a CrewBase decorator, None otherwise.
|
||||||
|
"""
|
||||||
|
if fullname in ("crewai.project.CrewBase", "crewai.project.crew_base.CrewBase"):
|
||||||
|
return self._crew_base_hook
|
||||||
|
return None
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _crew_base_hook(ctx: ClassDefContext) -> None:
|
||||||
|
"""Add injected attributes to @CrewBase decorated classes.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
ctx: Context for the class being decorated.
|
||||||
|
"""
|
||||||
|
any_type = AnyType(TypeOfAny.explicit)
|
||||||
|
str_type = ctx.api.named_type("builtins.str")
|
||||||
|
dict_type = ctx.api.named_type("builtins.dict", [str_type, any_type])
|
||||||
|
agents_config_var = Var("agents_config", dict_type)
|
||||||
|
agents_config_var.info = ctx.cls.info
|
||||||
|
agents_config_var._fullname = f"{ctx.cls.info.fullname}.agents_config"
|
||||||
|
ctx.cls.info.names["agents_config"] = SymbolTableNode(MDEF, agents_config_var)
|
||||||
|
tasks_config_var = Var("tasks_config", dict_type)
|
||||||
|
tasks_config_var.info = ctx.cls.info
|
||||||
|
tasks_config_var._fullname = f"{ctx.cls.info.fullname}.tasks_config"
|
||||||
|
ctx.cls.info.names["tasks_config"] = SymbolTableNode(MDEF, tasks_config_var)
|
||||||
|
|
||||||
|
|
||||||
|
def plugin(_: str) -> type[Plugin]:
|
||||||
|
"""Entry point for mypy plugin.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
_: Mypy version string.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Plugin class.
|
||||||
|
"""
|
||||||
|
return CrewAIPlugin
|
||||||
@@ -20,7 +20,7 @@ from typing_extensions import Self
|
|||||||
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from crewai import Agent, Task
|
from crewai import Agent, Crew, Task
|
||||||
from crewai.crews.crew_output import CrewOutput
|
from crewai.crews.crew_output import CrewOutput
|
||||||
from crewai.tools import BaseTool
|
from crewai.tools import BaseTool
|
||||||
|
|
||||||
@@ -129,6 +129,7 @@ class CrewClass(Protocol):
|
|||||||
_map_agent_variables: Callable[..., None]
|
_map_agent_variables: Callable[..., None]
|
||||||
map_all_task_variables: Callable[..., None]
|
map_all_task_variables: Callable[..., None]
|
||||||
_map_task_variables: Callable[..., None]
|
_map_task_variables: Callable[..., None]
|
||||||
|
crew: Callable[..., Crew]
|
||||||
|
|
||||||
|
|
||||||
class DecoratedMethod(Generic[P, R]):
|
class DecoratedMethod(Generic[P, R]):
|
||||||
|
|||||||
@@ -124,7 +124,7 @@ exclude = [
|
|||||||
"lib/crewai-tools/tests/",
|
"lib/crewai-tools/tests/",
|
||||||
"lib/crewai/src/crewai/experimental/a2a"
|
"lib/crewai/src/crewai/experimental/a2a"
|
||||||
]
|
]
|
||||||
plugins = ["pydantic.mypy"]
|
plugins = ["pydantic.mypy", "crewai.mypy"]
|
||||||
|
|
||||||
|
|
||||||
[tool.bandit]
|
[tool.bandit]
|
||||||
|
|||||||
Reference in New Issue
Block a user