Coverage for src/quickmcp/decorators.py: 100%

1""" 

2QuickMCP Decorators - Standalone decorators for MCP components 

3""" 

4 

5from typing import Callable, Optional, Dict, Any, List 

6from functools import wraps 

7import inspect 

8import logging 

9 

10logger = logging.getLogger(__name__) 

11 

12# Global registry for decorated functions (before server is created) 

13_pending_tools: List[Dict[str, Any]] = [] 

14_pending_resources: List[Dict[str, Any]] = [] 

15_pending_prompts: List[Dict[str, Any]] = [] 

16 

17 

18def tool( 

19 name: Optional[str] = None, 

20 description: Optional[str] = None, 

21 schema: Optional[Dict[str, Any]] = None, 

22) -> Callable: 

23 """ 

24 Standalone decorator to mark a function as an MCP tool. 

25  

26 Can be used before creating a server instance: 

27  

28 ```python 

29 @tool() 

30 def calculate(a: int, b: int) -> int: 

31 return a + b 

32  

33 # Later... 

34 server = QuickMCPServer("my-server") 

35 server.register_decorated() # Registers all decorated functions 

36 ``` 

37 """ 

38 def decorator(func: Callable) -> Callable: 

39 tool_name = name or func.__name__ 

40 tool_desc = description or (func.__doc__ or "").strip() 

41 

42 # Store metadata for later registration 

43 _pending_tools.append({ 

44 "function": func, 

45 "name": tool_name, 

46 "description": tool_desc, 

47 "schema": schema, 

48 }) 

49 

50 # Mark the function 

51 func._mcp_tool = True 

52 func._mcp_metadata = { 

53 "name": tool_name, 

54 "description": tool_desc, 

55 "schema": schema, 

56 } 

57 

58 @wraps(func) 

59 def wrapper(*args, **kwargs): 

60 return func(*args, **kwargs) 

61 

62 return wrapper 

63 

64 return decorator 

65 

66 

67def resource( 

68 uri_template: str, 

69 name: Optional[str] = None, 

70 description: Optional[str] = None, 

71 mime_type: str = "text/plain", 

72) -> Callable: 

73 """ 

74 Standalone decorator to mark a function as an MCP resource. 

75  

76 ```python 

77 @resource("config://{key}") 

78 def get_config(key: str) -> str: 

79 return config_data[key] 

80 ``` 

81 """ 

82 def decorator(func: Callable) -> Callable: 

83 resource_name = name or func.__name__ 

84 resource_desc = description or (func.__doc__ or "").strip() 

85 

86 # Store metadata for later registration 

87 _pending_resources.append({ 

88 "function": func, 

89 "uri_template": uri_template, 

90 "name": resource_name, 

91 "description": resource_desc, 

92 "mime_type": mime_type, 

93 }) 

94 

95 # Mark the function 

96 func._mcp_resource = True 

97 func._mcp_metadata = { 

98 "uri_template": uri_template, 

99 "name": resource_name, 

100 "description": resource_desc, 

101 "mime_type": mime_type, 

102 } 

103 

104 @wraps(func) 

105 def wrapper(*args, **kwargs): 

106 return func(*args, **kwargs) 

107 

108 return wrapper 

109 

110 return decorator 

111 

112 

113def prompt( 

114 name: Optional[str] = None, 

115 description: Optional[str] = None, 

116 arguments: Optional[List[Dict[str, Any]]] = None, 

117) -> Callable: 

118 """ 

119 Standalone decorator to mark a function as an MCP prompt template. 

120  

121 ```python 

122 @prompt() 

123 def code_review(language: str, code: str) -> str: 

124 return f"Review this {language} code: {code}" 

125 ``` 

126 """ 

127 def decorator(func: Callable) -> Callable: 

128 prompt_name = name or func.__name__ 

129 prompt_desc = description or (func.__doc__ or "").strip() 

130 

131 # Store metadata for later registration 

132 _pending_prompts.append({ 

133 "function": func, 

134 "name": prompt_name, 

135 "description": prompt_desc, 

136 "arguments": arguments, 

137 }) 

138 

139 # Mark the function 

140 func._mcp_prompt = True 

141 func._mcp_metadata = { 

142 "name": prompt_name, 

143 "description": prompt_desc, 

144 "arguments": arguments, 

145 } 

146 

147 @wraps(func) 

148 def wrapper(*args, **kwargs): 

149 return func(*args, **kwargs) 

150 

151 return wrapper 

152 

153 return decorator 

154 

155 

156def get_pending_tools() -> List[Dict[str, Any]]: 

157 """Get all pending tools waiting for registration.""" 

158 return _pending_tools.copy() 

159 

160 

161def get_pending_resources() -> List[Dict[str, Any]]: 

162 """Get all pending resources waiting for registration.""" 

163 return _pending_resources.copy() 

164 

165 

166def get_pending_prompts() -> List[Dict[str, Any]]: 

167 """Get all pending prompts waiting for registration.""" 

168 return _pending_prompts.copy() 

169 

170 

171def clear_pending() -> None: 

172 """Clear all pending registrations.""" 

173 _pending_tools.clear() 

174 _pending_resources.clear() 

175 _pending_prompts.clear() 

176 

177 

178def auto_discover_mcp_components(module) -> Dict[str, List[Callable]]: 

179 """ 

180 Auto-discover MCP components in a module. 

181  

182 Args: 

183 module: Python module to scan 

184  

185 Returns: 

186 Dictionary with discovered tools, resources, and prompts 

187 """ 

188 discovered = { 

189 "tools": [], 

190 "resources": [], 

191 "prompts": [], 

192 } 

193 

194 for name, obj in inspect.getmembers(module): 

195 if callable(obj): 

196 if hasattr(obj, "_mcp_tool"): 

197 discovered["tools"].append(obj) 

198 elif hasattr(obj, "_mcp_resource"): 

199 discovered["resources"].append(obj) 

200 elif hasattr(obj, "_mcp_prompt"): 

201 discovered["prompts"].append(obj) 

202 

203 return discovered