plain.models 0.49.1__py3-none-any.whl → 0.50.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (105) hide show
  1. plain/models/CHANGELOG.md +23 -0
  2. plain/models/aggregates.py +42 -19
  3. plain/models/backends/base/base.py +125 -105
  4. plain/models/backends/base/client.py +11 -3
  5. plain/models/backends/base/creation.py +22 -12
  6. plain/models/backends/base/features.py +10 -4
  7. plain/models/backends/base/introspection.py +29 -16
  8. plain/models/backends/base/operations.py +187 -91
  9. plain/models/backends/base/schema.py +267 -165
  10. plain/models/backends/base/validation.py +12 -3
  11. plain/models/backends/ddl_references.py +85 -43
  12. plain/models/backends/mysql/base.py +29 -26
  13. plain/models/backends/mysql/client.py +7 -2
  14. plain/models/backends/mysql/compiler.py +12 -3
  15. plain/models/backends/mysql/creation.py +5 -2
  16. plain/models/backends/mysql/features.py +24 -22
  17. plain/models/backends/mysql/introspection.py +22 -13
  18. plain/models/backends/mysql/operations.py +106 -39
  19. plain/models/backends/mysql/schema.py +48 -24
  20. plain/models/backends/mysql/validation.py +13 -6
  21. plain/models/backends/postgresql/base.py +41 -34
  22. plain/models/backends/postgresql/client.py +7 -2
  23. plain/models/backends/postgresql/creation.py +10 -5
  24. plain/models/backends/postgresql/introspection.py +15 -8
  25. plain/models/backends/postgresql/operations.py +109 -42
  26. plain/models/backends/postgresql/schema.py +85 -46
  27. plain/models/backends/sqlite3/_functions.py +151 -115
  28. plain/models/backends/sqlite3/base.py +37 -23
  29. plain/models/backends/sqlite3/client.py +7 -1
  30. plain/models/backends/sqlite3/creation.py +9 -5
  31. plain/models/backends/sqlite3/features.py +5 -3
  32. plain/models/backends/sqlite3/introspection.py +32 -16
  33. plain/models/backends/sqlite3/operations.py +125 -42
  34. plain/models/backends/sqlite3/schema.py +82 -58
  35. plain/models/backends/utils.py +52 -29
  36. plain/models/backups/cli.py +8 -6
  37. plain/models/backups/clients.py +16 -7
  38. plain/models/backups/core.py +24 -13
  39. plain/models/base.py +113 -74
  40. plain/models/cli.py +94 -63
  41. plain/models/config.py +1 -1
  42. plain/models/connections.py +23 -7
  43. plain/models/constraints.py +65 -47
  44. plain/models/database_url.py +1 -1
  45. plain/models/db.py +6 -2
  46. plain/models/deletion.py +66 -43
  47. plain/models/entrypoints.py +1 -1
  48. plain/models/enums.py +22 -11
  49. plain/models/exceptions.py +23 -8
  50. plain/models/expressions.py +440 -257
  51. plain/models/fields/__init__.py +253 -202
  52. plain/models/fields/json.py +120 -54
  53. plain/models/fields/mixins.py +12 -8
  54. plain/models/fields/related.py +284 -252
  55. plain/models/fields/related_descriptors.py +34 -25
  56. plain/models/fields/related_lookups.py +23 -11
  57. plain/models/fields/related_managers.py +81 -47
  58. plain/models/fields/reverse_related.py +58 -55
  59. plain/models/forms.py +89 -63
  60. plain/models/functions/comparison.py +71 -18
  61. plain/models/functions/datetime.py +79 -29
  62. plain/models/functions/math.py +43 -10
  63. plain/models/functions/mixins.py +24 -7
  64. plain/models/functions/text.py +104 -25
  65. plain/models/functions/window.py +12 -6
  66. plain/models/indexes.py +52 -28
  67. plain/models/lookups.py +228 -153
  68. plain/models/migrations/autodetector.py +86 -43
  69. plain/models/migrations/exceptions.py +7 -3
  70. plain/models/migrations/executor.py +33 -7
  71. plain/models/migrations/graph.py +79 -50
  72. plain/models/migrations/loader.py +45 -22
  73. plain/models/migrations/migration.py +23 -18
  74. plain/models/migrations/operations/base.py +37 -19
  75. plain/models/migrations/operations/fields.py +89 -42
  76. plain/models/migrations/operations/models.py +245 -143
  77. plain/models/migrations/operations/special.py +82 -25
  78. plain/models/migrations/optimizer.py +7 -2
  79. plain/models/migrations/questioner.py +58 -31
  80. plain/models/migrations/recorder.py +18 -11
  81. plain/models/migrations/serializer.py +50 -39
  82. plain/models/migrations/state.py +220 -133
  83. plain/models/migrations/utils.py +29 -13
  84. plain/models/migrations/writer.py +17 -14
  85. plain/models/options.py +63 -56
  86. plain/models/otel.py +16 -6
  87. plain/models/preflight.py +35 -12
  88. plain/models/query.py +323 -228
  89. plain/models/query_utils.py +93 -58
  90. plain/models/registry.py +34 -16
  91. plain/models/sql/compiler.py +146 -97
  92. plain/models/sql/datastructures.py +38 -25
  93. plain/models/sql/query.py +255 -169
  94. plain/models/sql/subqueries.py +32 -21
  95. plain/models/sql/where.py +54 -29
  96. plain/models/test/pytest.py +15 -11
  97. plain/models/test/utils.py +4 -2
  98. plain/models/transaction.py +20 -7
  99. plain/models/utils.py +13 -5
  100. {plain_models-0.49.1.dist-info → plain_models-0.50.0.dist-info}/METADATA +1 -1
  101. plain_models-0.50.0.dist-info/RECORD +122 -0
  102. plain_models-0.49.1.dist-info/RECORD +0 -122
  103. {plain_models-0.49.1.dist-info → plain_models-0.50.0.dist-info}/WHEEL +0 -0
  104. {plain_models-0.49.1.dist-info → plain_models-0.50.0.dist-info}/entry_points.txt +0 -0
  105. {plain_models-0.49.1.dist-info → plain_models-0.50.0.dist-info}/licenses/LICENSE +0 -0
@@ -1,6 +1,9 @@
1
+ from __future__ import annotations
2
+
1
3
  import functools
2
4
  import re
3
5
  from graphlib import TopologicalSorter
6
+ from typing import TYPE_CHECKING, Any
4
7
 
5
8
  from plain import models
6
9
  from plain.models.migrations import operations
@@ -15,6 +18,12 @@ from plain.models.migrations.utils import (
15
18
  )
16
19
  from plain.runtime import settings
17
20
 
21
+ if TYPE_CHECKING:
22
+ from plain.models.fields import Field
23
+ from plain.models.migrations.graph import MigrationGraph
24
+ from plain.models.migrations.operations.base import Operation
25
+ from plain.models.migrations.state import ProjectState
26
+
18
27
 
19
28
  class MigrationAutodetector:
20
29
  """
@@ -29,15 +38,24 @@ class MigrationAutodetector:
29
38
  if it wishes, with the caveat that it may not always be possible.
30
39
  """
31
40
 
32
- def __init__(self, from_state, to_state, questioner=None):
41
+ def __init__(
42
+ self,
43
+ from_state: ProjectState,
44
+ to_state: ProjectState,
45
+ questioner: MigrationQuestioner | None = None,
46
+ ):
33
47
  self.from_state = from_state
34
48
  self.to_state = to_state
35
49
  self.questioner = questioner or MigrationQuestioner()
36
50
  self.existing_packages = {app for app, model in from_state.models}
37
51
 
38
52
  def changes(
39
- self, graph, trim_to_packages=None, convert_packages=None, migration_name=None
40
- ):
53
+ self,
54
+ graph: MigrationGraph,
55
+ trim_to_packages: set[str] | None = None,
56
+ convert_packages: set[str] | None = None,
57
+ migration_name: str | None = None,
58
+ ) -> dict[str, list[Migration]]:
41
59
  """
42
60
  Main entry point to produce a list of applicable changes.
43
61
  Take a graph to base names on and an optional set of packages
@@ -49,7 +67,7 @@ class MigrationAutodetector:
49
67
  changes = self._trim_to_packages(changes, trim_to_packages)
50
68
  return changes
51
69
 
52
- def deep_deconstruct(self, obj):
70
+ def deep_deconstruct(self, obj: Any) -> Any:
53
71
  """
54
72
  Recursive deconstruction for a field and its arguments.
55
73
  Used for full comparison for rename/alter; sometimes a single-level
@@ -87,7 +105,7 @@ class MigrationAutodetector:
87
105
  else:
88
106
  return obj
89
107
 
90
- def only_relation_agnostic_fields(self, fields):
108
+ def only_relation_agnostic_fields(self, fields: dict[str, Field]) -> list[Any]:
91
109
  """
92
110
  Return a definition of the fields that ignores field names and
93
111
  what related fields actually relate to. Used for detecting renames (as
@@ -101,7 +119,11 @@ class MigrationAutodetector:
101
119
  fields_def.append(deconstruction)
102
120
  return fields_def
103
121
 
104
- def _detect_changes(self, convert_packages=None, graph=None):
122
+ def _detect_changes(
123
+ self,
124
+ convert_packages: set[str] | None = None,
125
+ graph: MigrationGraph | None = None,
126
+ ) -> dict[str, list[Migration]]:
105
127
  """
106
128
  Return a dict of migration plans which will achieve the
107
129
  change from from_state to to_state. The dict has app labels
@@ -184,7 +206,7 @@ class MigrationAutodetector:
184
206
 
185
207
  return self.migrations
186
208
 
187
- def _prepare_field_lists(self):
209
+ def _prepare_field_lists(self) -> None:
188
210
  """
189
211
  Prepare field lists and a list of the fields that used through models
190
212
  in the old state so dependencies can be made from the through model
@@ -206,7 +228,7 @@ class MigrationAutodetector:
206
228
  for field_name in self.to_state.models[package_label, model_name].fields
207
229
  }
208
230
 
209
- def _generate_through_model_map(self):
231
+ def _generate_through_model_map(self) -> None:
210
232
  """Through model map generation."""
211
233
  for package_label, model_name in sorted(self.old_model_keys):
212
234
  old_model_name = self.renamed_models.get(
@@ -227,7 +249,9 @@ class MigrationAutodetector:
227
249
  )
228
250
 
229
251
  @staticmethod
230
- def _resolve_dependency(dependency):
252
+ def _resolve_dependency(
253
+ dependency: tuple[str, ...],
254
+ ) -> tuple[tuple[str, ...], bool]:
231
255
  """
232
256
  Return the resolved dependency and a boolean denoting whether or not
233
257
  it was a settings dependency.
@@ -241,7 +265,7 @@ class MigrationAutodetector:
241
265
  2:
242
266
  ], True
243
267
 
244
- def _build_migration_list(self, graph=None):
268
+ def _build_migration_list(self, graph: MigrationGraph | None = None) -> None:
245
269
  """
246
270
  Chop the lists of operations up into migrations with dependencies on
247
271
  each other. Do this by going through an app's list of operations until
@@ -355,7 +379,7 @@ class MigrationAutodetector:
355
379
  )
356
380
  num_ops = new_num_ops
357
381
 
358
- def _sort_migrations(self):
382
+ def _sort_migrations(self) -> None:
359
383
  """
360
384
  Reorder to make things possible. Reordering may be needed so FKs work
361
385
  nicely inside the same app.
@@ -373,7 +397,7 @@ class MigrationAutodetector:
373
397
  ts.add(op, *(x for x in ops if self.check_dependency(x, dep)))
374
398
  self.generated_operations[package_label] = list(ts.static_order())
375
399
 
376
- def _optimize_migrations(self):
400
+ def _optimize_migrations(self) -> None:
377
401
  # Add in internal dependencies among the migrations
378
402
  for package_label, migrations in self.migrations.items():
379
403
  for m1, m2 in zip(migrations, migrations[1:]):
@@ -391,7 +415,9 @@ class MigrationAutodetector:
391
415
  migration.operations, package_label
392
416
  )
393
417
 
394
- def check_dependency(self, operation, dependency):
418
+ def check_dependency(
419
+ self, operation: Operation, dependency: tuple[str, ...]
420
+ ) -> bool:
395
421
  """
396
422
  Return True if the given operation depends on the given dependency,
397
423
  False otherwise.
@@ -438,17 +464,21 @@ class MigrationAutodetector:
438
464
  raise ValueError(f"Can't handle dependency {dependency!r}")
439
465
 
440
466
  def add_operation(
441
- self, package_label, operation, dependencies=None, beginning=False
442
- ):
467
+ self,
468
+ package_label: str,
469
+ operation: Operation,
470
+ dependencies: list[tuple[str, ...]] | None = None,
471
+ beginning: bool = False,
472
+ ) -> None:
443
473
  # Dependencies are
444
474
  # (package_label, model_name, field_name, create/delete as True/False)
445
- operation._auto_deps = dependencies or []
475
+ operation._auto_deps = dependencies or [] # type: ignore[attr-defined]
446
476
  if beginning:
447
477
  self.generated_operations.setdefault(package_label, []).insert(0, operation)
448
478
  else:
449
479
  self.generated_operations.setdefault(package_label, []).append(operation)
450
480
 
451
- def generate_renamed_models(self):
481
+ def generate_renamed_models(self) -> None:
452
482
  """
453
483
  Find any renamed models, generate the operations for them, and remove
454
484
  the old entry from the model lists. Must be run before other
@@ -513,7 +543,7 @@ class MigrationAutodetector:
513
543
  self.old_model_keys.add((package_label, model_name))
514
544
  break
515
545
 
516
- def generate_created_models(self):
546
+ def generate_created_models(self) -> None:
517
547
  """
518
548
  Find all new models and make create
519
549
  operations for them as well as separate operations to create any
@@ -649,7 +679,7 @@ class MigrationAutodetector:
649
679
  dependencies=related_dependencies,
650
680
  )
651
681
 
652
- def generate_deleted_models(self):
682
+ def generate_deleted_models(self) -> None:
653
683
  """
654
684
  Find all deleted models and make delete
655
685
  operations for them as well as separate operations to delete any
@@ -724,7 +754,7 @@ class MigrationAutodetector:
724
754
  dependencies=list(set(dependencies)),
725
755
  )
726
756
 
727
- def create_renamed_fields(self):
757
+ def create_renamed_fields(self) -> None:
728
758
  """Work out renamed fields."""
729
759
  self.renamed_operations = []
730
760
  old_field_keys = self.old_field_keys.copy()
@@ -786,7 +816,7 @@ class MigrationAutodetector:
786
816
  ] = rem_field_name
787
817
  break
788
818
 
789
- def generate_renamed_fields(self):
819
+ def generate_renamed_fields(self) -> None:
790
820
  """Generate RenameField operations."""
791
821
  for (
792
822
  rem_package_label,
@@ -825,14 +855,16 @@ class MigrationAutodetector:
825
855
  )
826
856
  self.old_field_keys.add((package_label, model_name, field_name))
827
857
 
828
- def generate_added_fields(self):
858
+ def generate_added_fields(self) -> None:
829
859
  """Make AddField operations."""
830
860
  for package_label, model_name, field_name in sorted(
831
861
  self.new_field_keys - self.old_field_keys
832
862
  ):
833
863
  self._generate_added_field(package_label, model_name, field_name)
834
864
 
835
- def _generate_added_field(self, package_label, model_name, field_name):
865
+ def _generate_added_field(
866
+ self, package_label: str, model_name: str, field_name: str
867
+ ) -> None:
836
868
  field = self.to_state.models[package_label, model_name].get_field(field_name)
837
869
  # Adding a field always depends at least on its removal.
838
870
  dependencies = [(package_label, model_name, field_name, False)]
@@ -883,14 +915,16 @@ class MigrationAutodetector:
883
915
  dependencies=dependencies,
884
916
  )
885
917
 
886
- def generate_removed_fields(self):
918
+ def generate_removed_fields(self) -> None:
887
919
  """Make RemoveField operations."""
888
920
  for package_label, model_name, field_name in sorted(
889
921
  self.old_field_keys - self.new_field_keys
890
922
  ):
891
923
  self._generate_removed_field(package_label, model_name, field_name)
892
924
 
893
- def _generate_removed_field(self, package_label, model_name, field_name):
925
+ def _generate_removed_field(
926
+ self, package_label: str, model_name: str, field_name: str
927
+ ) -> None:
894
928
  self.add_operation(
895
929
  package_label,
896
930
  operations.RemoveField(
@@ -899,7 +933,7 @@ class MigrationAutodetector:
899
933
  ),
900
934
  )
901
935
 
902
- def generate_altered_fields(self):
936
+ def generate_altered_fields(self) -> None:
903
937
  """
904
938
  Make AlterField operations, or possibly RemovedField/AddField if alter
905
939
  isn't possible.
@@ -995,7 +1029,7 @@ class MigrationAutodetector:
995
1029
  self._generate_removed_field(package_label, model_name, field_name)
996
1030
  self._generate_added_field(package_label, model_name, field_name)
997
1031
 
998
- def create_altered_indexes(self):
1032
+ def create_altered_indexes(self) -> None:
999
1033
  option_name = operations.AddIndex.option_name
1000
1034
 
1001
1035
  for package_label, model_name in sorted(self.kept_model_keys):
@@ -1047,7 +1081,7 @@ class MigrationAutodetector:
1047
1081
  }
1048
1082
  )
1049
1083
 
1050
- def generate_added_indexes(self):
1084
+ def generate_added_indexes(self) -> None:
1051
1085
  for (package_label, model_name), alt_indexes in self.altered_indexes.items():
1052
1086
  dependencies = self._get_dependencies_for_model(package_label, model_name)
1053
1087
  for index in alt_indexes["added_indexes"]:
@@ -1060,7 +1094,7 @@ class MigrationAutodetector:
1060
1094
  dependencies=dependencies,
1061
1095
  )
1062
1096
 
1063
- def generate_removed_indexes(self):
1097
+ def generate_removed_indexes(self) -> None:
1064
1098
  for (package_label, model_name), alt_indexes in self.altered_indexes.items():
1065
1099
  for index in alt_indexes["removed_indexes"]:
1066
1100
  self.add_operation(
@@ -1071,7 +1105,7 @@ class MigrationAutodetector:
1071
1105
  ),
1072
1106
  )
1073
1107
 
1074
- def generate_renamed_indexes(self):
1108
+ def generate_renamed_indexes(self) -> None:
1075
1109
  for (package_label, model_name), alt_indexes in self.altered_indexes.items():
1076
1110
  for old_index_name, new_index_name, old_fields in alt_indexes[
1077
1111
  "renamed_indexes"
@@ -1086,7 +1120,7 @@ class MigrationAutodetector:
1086
1120
  ),
1087
1121
  )
1088
1122
 
1089
- def create_altered_constraints(self):
1123
+ def create_altered_constraints(self) -> None:
1090
1124
  option_name = operations.AddConstraint.option_name
1091
1125
  for package_label, model_name in sorted(self.kept_model_keys):
1092
1126
  old_model_name = self.renamed_models.get(
@@ -1109,7 +1143,7 @@ class MigrationAutodetector:
1109
1143
  }
1110
1144
  )
1111
1145
 
1112
- def generate_added_constraints(self):
1146
+ def generate_added_constraints(self) -> None:
1113
1147
  for (
1114
1148
  package_label,
1115
1149
  model_name,
@@ -1125,7 +1159,7 @@ class MigrationAutodetector:
1125
1159
  dependencies=dependencies,
1126
1160
  )
1127
1161
 
1128
- def generate_removed_constraints(self):
1162
+ def generate_removed_constraints(self) -> None:
1129
1163
  for (
1130
1164
  package_label,
1131
1165
  model_name,
@@ -1141,8 +1175,8 @@ class MigrationAutodetector:
1141
1175
 
1142
1176
  @staticmethod
1143
1177
  def _get_dependencies_for_foreign_key(
1144
- package_label, model_name, field, project_state
1145
- ):
1178
+ package_label: str, model_name: str, field: Field, project_state: ProjectState
1179
+ ) -> list[tuple[str, str, None, bool]]:
1146
1180
  remote_field_model = None
1147
1181
  if hasattr(field.remote_field, "model"):
1148
1182
  remote_field_model = field.remote_field.model
@@ -1163,7 +1197,7 @@ class MigrationAutodetector:
1163
1197
  dependencies = [(dep_package_label, dep_object_name, None, True)]
1164
1198
  if getattr(field.remote_field, "through", None):
1165
1199
  through_package_label, through_object_name = resolve_relation(
1166
- field.remote_field.through,
1200
+ field.remote_field.through, # type: ignore[attr-defined]
1167
1201
  package_label,
1168
1202
  model_name,
1169
1203
  )
@@ -1172,7 +1206,9 @@ class MigrationAutodetector:
1172
1206
  )
1173
1207
  return dependencies
1174
1208
 
1175
- def _get_dependencies_for_model(self, package_label, model_name):
1209
+ def _get_dependencies_for_model(
1210
+ self, package_label: str, model_name: str
1211
+ ) -> list[tuple[str, str, None, bool]]:
1176
1212
  """Return foreign key dependencies of the given model."""
1177
1213
  dependencies = []
1178
1214
  model_state = self.to_state.models[package_label, model_name]
@@ -1188,7 +1224,7 @@ class MigrationAutodetector:
1188
1224
  )
1189
1225
  return dependencies
1190
1226
 
1191
- def generate_altered_db_table(self):
1227
+ def generate_altered_db_table(self) -> None:
1192
1228
  for package_label, model_name in sorted(self.kept_model_keys):
1193
1229
  old_model_name = self.renamed_models.get(
1194
1230
  (package_label, model_name), model_name
@@ -1206,7 +1242,7 @@ class MigrationAutodetector:
1206
1242
  ),
1207
1243
  )
1208
1244
 
1209
- def generate_altered_db_table_comment(self):
1245
+ def generate_altered_db_table_comment(self) -> None:
1210
1246
  for package_label, model_name in sorted(self.kept_model_keys):
1211
1247
  old_model_name = self.renamed_models.get(
1212
1248
  (package_label, model_name), model_name
@@ -1225,7 +1261,7 @@ class MigrationAutodetector:
1225
1261
  ),
1226
1262
  )
1227
1263
 
1228
- def generate_altered_options(self):
1264
+ def generate_altered_options(self) -> None:
1229
1265
  """
1230
1266
  Work out if any non-schema-affecting options have changed and make an
1231
1267
  operation to represent them in state changes (in case Python code in
@@ -1256,7 +1292,12 @@ class MigrationAutodetector:
1256
1292
  ),
1257
1293
  )
1258
1294
 
1259
- def arrange_for_graph(self, changes, graph, migration_name=None):
1295
+ def arrange_for_graph(
1296
+ self,
1297
+ changes: dict[str, list[Migration]],
1298
+ graph: MigrationGraph,
1299
+ migration_name: str | None = None,
1300
+ ) -> dict[str, list[Migration]]:
1260
1301
  """
1261
1302
  Take a result from changes() and a MigrationGraph, and fix the names
1262
1303
  and dependencies of the changes so they extend the graph from the leaf
@@ -1311,7 +1352,9 @@ class MigrationAutodetector:
1311
1352
  ]
1312
1353
  return changes
1313
1354
 
1314
- def _trim_to_packages(self, changes, package_labels):
1355
+ def _trim_to_packages(
1356
+ self, changes: dict[str, list[Migration]], package_labels: set[str]
1357
+ ) -> dict[str, list[Migration]]:
1315
1358
  """
1316
1359
  Take changes from arrange_for_graph() and set of app labels, and return
1317
1360
  a modified set of changes which trims out as many migrations that are
@@ -1344,7 +1387,7 @@ class MigrationAutodetector:
1344
1387
  return changes
1345
1388
 
1346
1389
  @classmethod
1347
- def parse_number(cls, name):
1390
+ def parse_number(cls, name: str) -> int | None:
1348
1391
  """
1349
1392
  Given a migration name, try to extract a number from the beginning of
1350
1393
  it. For a squashed migration such as '0001_squashed_0004…', return the
@@ -1,3 +1,7 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Any
4
+
1
5
  from plain.models.db import DatabaseError
2
6
 
3
7
 
@@ -34,15 +38,15 @@ class InvalidBasesError(ValueError):
34
38
  class NodeNotFoundError(LookupError):
35
39
  """An attempt on a node is made that is not available in the graph."""
36
40
 
37
- def __init__(self, message, node, origin=None):
41
+ def __init__(self, message: str, node: Any, origin: Any = None) -> None:
38
42
  self.message = message
39
43
  self.origin = origin
40
44
  self.node = node
41
45
 
42
- def __str__(self):
46
+ def __str__(self) -> str:
43
47
  return self.message
44
48
 
45
- def __repr__(self):
49
+ def __repr__(self) -> str:
46
50
  return f"NodeNotFoundError({self.node!r})"
47
51
 
48
52
 
@@ -1,10 +1,18 @@
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Callable
1
4
  from contextlib import nullcontext
5
+ from typing import TYPE_CHECKING, Any
2
6
 
3
7
  from ..transaction import atomic
4
8
  from .loader import MigrationLoader
9
+ from .migration import Migration
5
10
  from .recorder import MigrationRecorder
6
11
  from .state import ProjectState
7
12
 
13
+ if TYPE_CHECKING:
14
+ from plain.models.backends.base.base import BaseDatabaseWrapper
15
+
8
16
 
9
17
  class MigrationExecutor:
10
18
  """
@@ -12,13 +20,19 @@ class MigrationExecutor:
12
20
  to a specified set of targets.
13
21
  """
14
22
 
15
- def __init__(self, connection, progress_callback=None):
23
+ def __init__(
24
+ self,
25
+ connection: BaseDatabaseWrapper,
26
+ progress_callback: Callable[..., Any] | None = None,
27
+ ) -> None:
16
28
  self.connection = connection
17
29
  self.loader = MigrationLoader(self.connection)
18
30
  self.recorder = MigrationRecorder(self.connection)
19
31
  self.progress_callback = progress_callback
20
32
 
21
- def migration_plan(self, targets, clean_start=False):
33
+ def migration_plan(
34
+ self, targets: list[tuple[str, str]], clean_start: bool = False
35
+ ) -> list[Migration]:
22
36
  """
23
37
  Given a set of targets, return a list of Migration instances.
24
38
  """
@@ -34,7 +48,9 @@ class MigrationExecutor:
34
48
  applied[migration] = self.loader.graph.nodes[migration]
35
49
  return plan
36
50
 
37
- def _create_project_state(self, with_applied_migrations=False):
51
+ def _create_project_state(
52
+ self, with_applied_migrations: bool = False
53
+ ) -> ProjectState:
38
54
  """
39
55
  Create a project state including all the applications without
40
56
  migrations and applied migrations if with_applied_migrations=True.
@@ -55,7 +71,14 @@ class MigrationExecutor:
55
71
  migration.mutate_state(state, preserve=False)
56
72
  return state
57
73
 
58
- def migrate(self, targets, plan=None, state=None, fake=False, atomic_batch=False):
74
+ def migrate(
75
+ self,
76
+ targets: list[tuple[str, str]],
77
+ plan: list[Migration] | None = None,
78
+ state: ProjectState | None = None,
79
+ fake: bool = False,
80
+ atomic_batch: bool = False,
81
+ ) -> ProjectState:
59
82
  """
60
83
  Migrate the database up to the given targets.
61
84
 
@@ -113,9 +136,12 @@ class MigrationExecutor:
113
136
 
114
137
  self.check_replacements()
115
138
 
139
+ assert state is not None
116
140
  return state
117
141
 
118
- def apply_migration(self, state, migration, fake=False):
142
+ def apply_migration(
143
+ self, state: ProjectState, migration: Migration, fake: bool = False
144
+ ) -> ProjectState:
119
145
  """Run a migration forwards."""
120
146
  migration_recorded = False
121
147
  if self.progress_callback:
@@ -136,7 +162,7 @@ class MigrationExecutor:
136
162
  self.progress_callback("apply_success", migration, fake)
137
163
  return state
138
164
 
139
- def record_migration(self, migration):
165
+ def record_migration(self, migration: Migration) -> None:
140
166
  # For replacement migrations, record individual statuses
141
167
  if migration.replaces:
142
168
  for package_label, name in migration.replaces:
@@ -144,7 +170,7 @@ class MigrationExecutor:
144
170
  else:
145
171
  self.recorder.record_applied(migration.package_label, migration.name)
146
172
 
147
- def check_replacements(self):
173
+ def check_replacements(self) -> None:
148
174
  """
149
175
  Mark replacement migrations applied if their replaced set all are.
150
176