Coverage for /Users/antonigmitruk/golf/src/golf/core/transformer.py: 0%

1"""Transform GolfMCP components into standalone FastMCP code. 

2 

3This module provides utilities for transforming GolfMCP's convention-based code 

4into explicit FastMCP component registrations. 

5""" 

6 

7import ast 

8from pathlib import Path 

9from typing import Any 

10 

11from golf.core.parser import ParsedComponent 

12 

13 

14class ImportTransformer(ast.NodeTransformer): 

15 """AST transformer for rewriting imports in component files.""" 

16 

17 def __init__( 

18 self, 

19 original_path: Path, 

20 target_path: Path, 

21 import_map: dict[str, str], 

22 project_root: Path, 

23 ) -> None: 

24 """Initialize the import transformer. 

25 

26 Args: 

27 original_path: Path to the original file 

28 target_path: Path to the target file 

29 import_map: Mapping of original module paths to generated paths 

30 project_root: Root path of the project 

31 """ 

32 self.original_path = original_path 

33 self.target_path = target_path 

34 self.import_map = import_map 

35 self.project_root = project_root 

36 

37 def visit_Import(self, node: ast.Import) -> Any: 

38 """Transform import statements.""" 

39 return node 

40 

41 def visit_ImportFrom(self, node: ast.ImportFrom) -> Any: 

42 """Transform import from statements.""" 

43 if node.module is None: 

44 return node 

45 

46 # Handle relative imports 

47 if node.level > 0: 

48 # Calculate the source module path 

49 source_dir = self.original_path.parent 

50 for _ in range(node.level - 1): 

51 source_dir = source_dir.parent 

52 

53 if node.module: 

54 # Handle imports like `from .helpers import utils` 

55 source_module = source_dir / node.module.replace(".", "/") 

56 else: 

57 # Handle imports like `from . import something` 

58 source_module = source_dir 

59 

60 try: 

61 # Check if this is a shared module import 

62 source_str = str(source_module.relative_to(self.project_root)) 

63 

64 # First, try direct module path match (e.g., "tools/weather/helpers") 

65 if source_str in self.import_map: 

66 new_module = self.import_map[source_str] 

67 return ast.ImportFrom(module=new_module, names=node.names, level=0) 

68 

69 # If direct match fails, try directory-based matching 

70 # This handles cases like `from . import common` where the import_map 

71 # has "tools/weather/common" but we're looking for "tools/weather" 

72 source_dir_str = str(source_dir.relative_to(self.project_root)) 

73 if source_dir_str in self.import_map: 

74 new_module = self.import_map[source_dir_str] 

75 if node.module: 

76 new_module = f"{new_module}.{node.module}" 

77 return ast.ImportFrom(module=new_module, names=node.names, level=0) 

78 

79 # Check for specific module imports within the directory 

80 for import_path, mapped_path in self.import_map.items(): 

81 # Handle cases where we import a specific module from a directory 

82 # e.g., `from .common import something` should match "tools/weather/common" 

83 if import_path.startswith(source_dir_str + "/") and node.module: 

84 module_name = import_path.replace(source_dir_str + "/", "") 

85 if module_name == node.module: 

86 return ast.ImportFrom(module=mapped_path, names=node.names, level=0) 

87 

88 except ValueError: 

89 # source_module is not relative to project_root, leave import unchanged 

90 pass 

91 

92 return node 

93 

94 

95def transform_component( 

96 component: ParsedComponent | None, 

97 output_file: Path, 

98 project_path: Path, 

99 import_map: dict[str, str], 

100 source_file: Path | None = None, 

101) -> str: 

102 """Transform a GolfMCP component into a standalone FastMCP component. 

103 

104 Args: 

105 component: Parsed component to transform (optional if source_file provided) 

106 output_file: Path to write the transformed component to 

107 project_path: Path to the project root 

108 import_map: Mapping of original module paths to generated paths 

109 source_file: Optional path to source file (for shared files) 

110 

111 Returns: 

112 Generated component code 

113 """ 

114 # Read the original file 

115 if source_file is not None: 

116 file_path = source_file 

117 elif component is not None: 

118 file_path = Path(component.file_path) 

119 else: 

120 raise ValueError("Either component or source_file must be provided") 

121 

122 with open(file_path) as f: 

123 source_code = f.read() 

124 

125 # Parse the source code into an AST 

126 tree = ast.parse(source_code) 

127 

128 # Transform imports 

129 transformer = ImportTransformer(file_path, output_file, import_map, project_path) 

130 tree = transformer.visit(tree) 

131 

132 # Get all imports and docstring 

133 imports = [] 

134 docstring = None 

135 

136 # Find the module docstring if present 

137 if ( 

138 len(tree.body) > 0 

139 and isinstance(tree.body[0], ast.Expr) 

140 and isinstance(tree.body[0].value, ast.Constant) 

141 and isinstance(tree.body[0].value.value, str) 

142 ): 

143 docstring = tree.body[0].value.value 

144 

145 # Find imports 

146 for node in tree.body: 

147 if isinstance(node, ast.Import | ast.ImportFrom): 

148 imports.append(node) 

149 

150 # Generate the transformed code 

151 transformed_imports = ast.unparse(ast.Module(body=imports, type_ignores=[])) 

152 

153 # Build full transformed code 

154 transformed_code = transformed_imports + "\n\n" 

155 

156 # Add docstring if present, using proper triple quotes for multi-line docstrings 

157 if docstring: 

158 # Check if docstring contains newlines 

159 if "\n" in docstring: 

160 # Use triple quotes for multi-line docstrings 

161 transformed_code += f'"""{docstring}"""\n\n' 

162 else: 

163 # Use single quotes for single-line docstrings 

164 transformed_code += f'"{docstring}"\n\n' 

165 

166 # Add the rest of the code except imports and the original docstring 

167 remaining_nodes = [] 

168 for node in tree.body: 

169 # Skip imports 

170 if isinstance(node, ast.Import | ast.ImportFrom): 

171 continue 

172 

173 # Skip the original docstring 

174 if isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant) and isinstance(node.value.value, str): 

175 continue 

176 

177 remaining_nodes.append(node) 

178 

179 remaining_code = ast.unparse(ast.Module(body=remaining_nodes, type_ignores=[])) 

180 transformed_code += remaining_code 

181 

182 # Ensure the directory exists 

183 output_file.parent.mkdir(parents=True, exist_ok=True) 

184 

185 # Write the transformed code to the output file 

186 with open(output_file, "w") as f: 

187 f.write(transformed_code) 

188 

189 return transformed_code