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

1 import os

2 import re

3 from importlib import import_module

4

5 from plain.models import migrations

6 from plain.models.migrations.loader import MigrationLoader

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

8 from plain.packages import packages

9 from plain.runtime import __version__

10 from plain.utils.inspect import get_func_args

11 from plain.utils.module_loading import module_dir

12 from plain.utils.timezone import now

13

14

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

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

286 MIGRATION_HEADER_TEMPLATE = """\

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

288

289"""

290

291

292 MIGRATION_TEMPLATE = """\

293%(migration_header)s%(imports)s

294

295 class Migration(migrations.Migration):

296%(replaces_str)s%(initial_str)s

297 dependencies = [

298%(dependencies)s\

299 ]

300

301 operations = [

302%(operations)s\

303 ]

304"""