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
« 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
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
15class OperationWriter:
16 def __init__(self, operation, indentation=2):
17 self.operation = operation
18 self.buff = []
19 self.indentation = indentation
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)
71 imports = set()
72 name, args, kwargs = self.operation.deconstruct()
73 operation_args = get_func_args(self.operation.__init__)
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}(")
84 self.indent()
86 for i, arg in enumerate(args):
87 arg_value = arg
88 arg_name = operation_args[i]
89 _write(arg_name, arg_value)
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)
98 self.unindent()
99 self.feed("),")
100 return self.render(), imports
102 def indent(self):
103 self.indentation += 1
105 def unindent(self):
106 self.indentation -= 1
108 def feed(self, line):
109 self.buff.append(" " * (self.indentation * 4) + line)
111 def render(self):
112 return "\n".join(self.buff)
115class MigrationWriter:
116 """
117 Take a Migration instance and is able to produce the contents
118 of the migration file from it.
119 """
121 def __init__(self, migration, include_header=True):
122 self.migration = migration
123 self.include_header = include_header
124 self.needs_manual_porting = False
126 def as_string(self):
127 """Return a string of the file contents."""
128 items = {
129 "replaces_str": "",
130 "initial_str": "",
131 }
133 imports = set()
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 ""
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 ""
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
164 imports.add("from plain.models import migrations")
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"] = ""
194 if self.migration.initial:
195 items["initial_str"] = "\n initial = True\n"
197 return MIGRATION_TEMPLATE % items
199 @property
200 def basedir(self):
201 migrations_package_name, _ = MigrationLoader.migrations_module(
202 self.migration.package_label
203 )
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 )
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
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)
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 )
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
263 return final_dir
265 @property
266 def filename(self):
267 return f"{self.migration.name}.py"
269 @property
270 def path(self):
271 return os.path.join(self.basedir, self.filename)
273 @classmethod
274 def serialize(cls, value):
275 return serializer_factory(value).serialize()
277 @classmethod
278 def register_serializer(cls, type_, serializer):
279 Serializer.register(type_, serializer)
281 @classmethod
282 def unregister_serializer(cls, type_):
283 Serializer.unregister(type_)
286MIGRATION_HEADER_TEMPLATE = """\
287# Generated by Plain %(version)s on %(timestamp)s
289"""
292MIGRATION_TEMPLATE = """\
293%(migration_header)s%(imports)s
295class Migration(migrations.Migration):
296%(replaces_str)s%(initial_str)s
297 dependencies = [
298%(dependencies)s\
299 ]
301 operations = [
302%(operations)s\
303 ]
304"""