Skip to content

Commit c594859

Browse files
feat: mypy plugin base
* feat: base mypy plugin with CrewBase * fix: add crew method to protocol
1 parent 2ee27ef commit c594859

File tree

3 files changed

+63
-2
lines changed

3 files changed

+63
-2
lines changed

lib/crewai/src/crewai/mypy.py

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
"""Mypy plugin for CrewAI decorator type checking.
2+
3+
This plugin informs mypy about attributes injected by the @CrewBase decorator.
4+
"""
5+
6+
from collections.abc import Callable
7+
8+
from mypy.nodes import MDEF, SymbolTableNode, Var
9+
from mypy.plugin import ClassDefContext, Plugin
10+
from mypy.types import AnyType, TypeOfAny
11+
12+
13+
class CrewAIPlugin(Plugin):
14+
"""Mypy plugin that handles @CrewBase decorator attribute injection."""
15+
16+
def get_class_decorator_hook(
17+
self, fullname: str
18+
) -> Callable[[ClassDefContext], None] | None:
19+
"""Return hook for class decorators.
20+
21+
Args:
22+
fullname: Fully qualified name of the decorator.
23+
24+
Returns:
25+
Hook function if this is a CrewBase decorator, None otherwise.
26+
"""
27+
if fullname in ("crewai.project.CrewBase", "crewai.project.crew_base.CrewBase"):
28+
return self._crew_base_hook
29+
return None
30+
31+
@staticmethod
32+
def _crew_base_hook(ctx: ClassDefContext) -> None:
33+
"""Add injected attributes to @CrewBase decorated classes.
34+
35+
Args:
36+
ctx: Context for the class being decorated.
37+
"""
38+
any_type = AnyType(TypeOfAny.explicit)
39+
str_type = ctx.api.named_type("builtins.str")
40+
dict_type = ctx.api.named_type("builtins.dict", [str_type, any_type])
41+
agents_config_var = Var("agents_config", dict_type)
42+
agents_config_var.info = ctx.cls.info
43+
agents_config_var._fullname = f"{ctx.cls.info.fullname}.agents_config"
44+
ctx.cls.info.names["agents_config"] = SymbolTableNode(MDEF, agents_config_var)
45+
tasks_config_var = Var("tasks_config", dict_type)
46+
tasks_config_var.info = ctx.cls.info
47+
tasks_config_var._fullname = f"{ctx.cls.info.fullname}.tasks_config"
48+
ctx.cls.info.names["tasks_config"] = SymbolTableNode(MDEF, tasks_config_var)
49+
50+
51+
def plugin(_: str) -> type[Plugin]:
52+
"""Entry point for mypy plugin.
53+
54+
Args:
55+
_: Mypy version string.
56+
57+
Returns:
58+
Plugin class.
59+
"""
60+
return CrewAIPlugin

lib/crewai/src/crewai/project/wrappers.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020

2121

2222
if TYPE_CHECKING:
23-
from crewai import Agent, Task
23+
from crewai import Agent, Crew, Task
2424
from crewai.crews.crew_output import CrewOutput
2525
from crewai.tools import BaseTool
2626

@@ -129,6 +129,7 @@ class CrewClass(Protocol):
129129
_map_agent_variables: Callable[..., None]
130130
map_all_task_variables: Callable[..., None]
131131
_map_task_variables: Callable[..., None]
132+
crew: Callable[..., Crew]
132133

133134

134135
class DecoratedMethod(Generic[P, R]):

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ exclude = [
124124
"lib/crewai-tools/tests/",
125125
"lib/crewai/src/crewai/experimental/a2a"
126126
]
127-
plugins = ["pydantic.mypy"]
127+
plugins = ["pydantic.mypy", "crewai.mypy"]
128128

129129

130130
[tool.bandit]

0 commit comments

Comments
 (0)