Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-models/plain/models/migrations/writer.py: 19%

180 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-23 11:16 -0600

1import os 

2import re 

3from importlib import import_module 

4 

5from plain.models import migrations 

6from plain.models.migrations.loader import MigrationLoader 

7from plain.models.migrations.serializer import Serializer, serializer_factory 

8from plain.packages import packages 

9from plain.runtime import __version__ 

10from plain.utils.inspect import get_func_args 

11from plain.utils.module_loading import module_dir 

12from plain.utils.timezone import now 

13 

14 

15class OperationWriter: 

16 def __init__(self, operation, indentation=2): 

17 self.operation = operation 

18 self.buff = [] 

19 self.indentation = indentation 

20 

21 def serialize(self): 

22 def _write(_arg_name, _arg_value): 

23 if _arg_name in self.operation.serialization_expand_args and isinstance( 

24 _arg_value, list | tuple | dict 

25 ): 

26 if isinstance(_arg_value, dict): 

27 self.feed(f"{_arg_name}={ ") 

28 self.indent() 

29 for key, value in _arg_value.items(): 

30 key_string, key_imports = MigrationWriter.serialize(key) 

31 arg_string, arg_imports = MigrationWriter.serialize(value) 

32 args = arg_string.splitlines() 

33 if len(args) > 1: 

34 self.feed(f"{key_string}: {args[0]}") 

35 for arg in args[1:-1]: 

36 self.feed(arg) 

37 self.feed(f"{args[-1]},") 

38 else: 

39 self.feed(f"{key_string}: {arg_string},") 

40 imports.update(key_imports) 

41 imports.update(arg_imports) 

42 self.unindent() 

43 self.feed("},") 

44 else: 

45 self.feed(f"{_arg_name}=[") 

46 self.indent() 

47 for item in _arg_value: 

48 arg_string, arg_imports = MigrationWriter.serialize(item) 

49 args = arg_string.splitlines() 

50 if len(args) > 1: 

51 for arg in args[:-1]: 

52 self.feed(arg) 

53 self.feed(f"{args[-1]},") 

54 else: 

55 self.feed(f"{arg_string},") 

56 imports.update(arg_imports) 

57 self.unindent() 

58 self.feed("],") 

59 else: 

60 arg_string, arg_imports = MigrationWriter.serialize(_arg_value) 

61 args = arg_string.splitlines() 

62 if len(args) > 1: 

63 self.feed(f"{_arg_name}={args[0]}") 

64 for arg in args[1:-1]: 

65 self.feed(arg) 

66 self.feed(f"{args[-1]},") 

67 else: 

68 self.feed(f"{_arg_name}={arg_string},") 

69 imports.update(arg_imports) 

70 

71 imports = set() 

72 name, args, kwargs = self.operation.deconstruct() 

73 operation_args = get_func_args(self.operation.__init__) 

74 

75 # See if this operation is in plain.models.migrations. If it is, 

76 # We can just use the fact we already have that imported, 

77 # otherwise, we need to add an import for the operation class. 

78 if getattr(migrations, name, None) == self.operation.__class__: 

79 self.feed(f"migrations.{name}(") 

80 else: 

81 imports.add(f"import {self.operation.__class__.__module__}") 

82 self.feed(f"{self.operation.__class__.__module__}.{name}(") 

83 

84 self.indent() 

85 

86 for i, arg in enumerate(args): 

87 arg_value = arg 

88 arg_name = operation_args[i] 

89 _write(arg_name, arg_value) 

90 

91 i = len(args) 

92 # Only iterate over remaining arguments 

93 for arg_name in operation_args[i:]: 

94 if arg_name in kwargs: # Don't sort to maintain signature order 

95 arg_value = kwargs[arg_name] 

96 _write(arg_name, arg_value) 

97 

98 self.unindent() 

99 self.feed("),") 

100 return self.render(), imports 

101 

102 def indent(self): 

103 self.indentation += 1 

104 

105 def unindent(self): 

106 self.indentation -= 1 

107 

108 def feed(self, line): 

109 self.buff.append(" " * (self.indentation * 4) + line) 

110 

111 def render(self): 

112 return "\n".join(self.buff) 

113 

114 

115class MigrationWriter: 

116 """ 

117 Take a Migration instance and is able to produce the contents 

118 of the migration file from it. 

119 """ 

120 

121 def __init__(self, migration, include_header=True): 

122 self.migration = migration 

123 self.include_header = include_header 

124 self.needs_manual_porting = False 

125 

126 def as_string(self): 

127 """Return a string of the file contents.""" 

128 items = { 

129 "replaces_str": "", 

130 "initial_str": "", 

131 } 

132 

133 imports = set() 

134 

135 # Deconstruct operations 

136 operations = [] 

137 for operation in self.migration.operations: 

138 operation_string, operation_imports = OperationWriter(operation).serialize() 

139 imports.update(operation_imports) 

140 operations.append(operation_string) 

141 items["operations"] = "\n".join(operations) + "\n" if operations else "" 

142 

143 # Format dependencies and write out swappable dependencies right 

144 dependencies = [] 

145 for dependency in self.migration.dependencies: 

146 if dependency[0] == "__setting__": 

147 dependencies.append( 

148 f" migrations.swappable_dependency(settings.{dependency[1]})," 

149 ) 

150 imports.add("from plain.runtime import settings") 

151 else: 

152 dependencies.append(f" {self.serialize(dependency)[0]},") 

153 items["dependencies"] = "\n".join(dependencies) + "\n" if dependencies else "" 

154 

155 # Format imports nicely, swapping imports of functions from migration files 

156 # for comments 

157 migration_imports = set() 

158 for line in list(imports): 

159 if re.match(r"^import (.*)\.\d+[^\s]*$", line): 

160 migration_imports.add(line.split("import")[1].strip()) 

161 imports.remove(line) 

162 self.needs_manual_porting = True 

163 

164 imports.add("from plain.models import migrations") 

165 

166 # Sort imports by the package / module to be imported (the part after 

167 # "from" in "from ... import ..." or after "import" in "import ..."). 

168 # First group the "import" statements, then "from ... import ...". 

169 sorted_imports = sorted( 

170 imports, key=lambda i: (i.split()[0] == "from", i.split()[1]) 

171 ) 

172 items["imports"] = "\n".join(sorted_imports) + "\n" if imports else "" 

173 if migration_imports: 

174 items["imports"] += ( 

175 "\n\n# Functions from the following migrations need manual " 

176 "copying.\n# Move them and any dependencies into this file, " 

177 "then update the\n# RunPython operations to refer to the local " 

178 "versions:\n# {}" 

179 ).format("\n# ".join(sorted(migration_imports))) 

180 # If there's a replaces, make a string for it 

181 if self.migration.replaces: 

182 items["replaces_str"] = ( 

183 f"\n replaces = {self.serialize(self.migration.replaces)[0]}\n" 

184 ) 

185 # Hinting that goes into comment 

186 if self.include_header: 

187 items["migration_header"] = MIGRATION_HEADER_TEMPLATE % { 

188 "version": __version__, 

189 "timestamp": now().strftime("%Y-%m-%d %H:%M"), 

190 } 

191 else: 

192 items["migration_header"] = "" 

193 

194 if self.migration.initial: 

195 items["initial_str"] = "\n initial = True\n" 

196 

197 return MIGRATION_TEMPLATE % items 

198 

199 @property 

200 def basedir(self): 

201 migrations_package_name, _ = MigrationLoader.migrations_module( 

202 self.migration.package_label 

203 ) 

204 

205 if migrations_package_name is None: 

206 raise ValueError( 

207 f"Plain can't create migrations for app '{self.migration.package_label}' because " 

208 "migrations have been disabled via the MIGRATION_MODULES " 

209 "setting." 

210 ) 

211 

212 # See if we can import the migrations module directly 

213 try: 

214 migrations_module = import_module(migrations_package_name) 

215 except ImportError: 

216 pass 

217 else: 

218 try: 

219 return module_dir(migrations_module) 

220 except ValueError: 

221 pass 

222 

223 # Alright, see if it's a direct submodule of the app 

224 package_config = packages.get_package_config(self.migration.package_label) 

225 ( 

226 maybe_package_name, 

227 _, 

228 migrations_package_basename, 

229 ) = migrations_package_name.rpartition(".") 

230 if package_config.name == maybe_package_name: 

231 return os.path.join(package_config.path, migrations_package_basename) 

232 

233 # In case of using MIGRATION_MODULES setting and the custom package 

234 # doesn't exist, create one, starting from an existing package 

235 existing_dirs, missing_dirs = migrations_package_name.split("."), [] 

236 while existing_dirs: 

237 missing_dirs.insert(0, existing_dirs.pop(-1)) 

238 try: 

239 base_module = import_module(".".join(existing_dirs)) 

240 except (ImportError, ValueError): 

241 continue 

242 else: 

243 try: 

244 base_dir = module_dir(base_module) 

245 except ValueError: 

246 continue 

247 else: 

248 break 

249 else: 

250 raise ValueError( 

251 "Could not locate an appropriate location to create " 

252 f"migrations package {migrations_package_name}. Make sure the toplevel " 

253 "package exists and can be imported." 

254 ) 

255 

256 final_dir = os.path.join(base_dir, *missing_dirs) 

257 os.makedirs(final_dir, exist_ok=True) 

258 for missing_dir in missing_dirs: 

259 base_dir = os.path.join(base_dir, missing_dir) 

260 with open(os.path.join(base_dir, "__init__.py"), "w"): 

261 pass 

262 

263 return final_dir 

264 

265 @property 

266 def filename(self): 

267 return f"{self.migration.name}.py" 

268 

269 @property 

270 def path(self): 

271 return os.path.join(self.basedir, self.filename) 

272 

273 @classmethod 

274 def serialize(cls, value): 

275 return serializer_factory(value).serialize() 

276 

277 @classmethod 

278 def register_serializer(cls, type_, serializer): 

279 Serializer.register(type_, serializer) 

280 

281 @classmethod 

282 def unregister_serializer(cls, type_): 

283 Serializer.unregister(type_) 

284 

285 

286MIGRATION_HEADER_TEMPLATE = """\ 

287# Generated by Plain %(version)s on %(timestamp)s 

288 

289""" 

290 

291 

292MIGRATION_TEMPLATE = """\ 

293%(migration_header)s%(imports)s 

294 

295class Migration(migrations.Migration): 

296%(replaces_str)s%(initial_str)s 

297 dependencies = [ 

298%(dependencies)s\ 

299 ] 

300 

301 operations = [ 

302%(operations)s\ 

303 ] 

304"""