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

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

2

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

4 into explicit FastMCP component registrations.

5"""

6

7 import ast

8 from pathlib import Path

9 from typing import Any

10

11 from golf.core.parser import ParsedComponent

12

13

14 class 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

95 def 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