From 7b31bbbf88175744698db1274078c20c5feb60b6 Mon Sep 17 00:00:00 2001 From: Waket Zheng Date: Sat, 16 May 2026 01:46:14 +0800 Subject: [PATCH 1/2] refactor: sorted imports in migration file --- tests/migrations/test_writer.py | 41 +++++++++++++-------------------- tortoise/migrations/writer.py | 21 ++++++++++------- 2 files changed, 29 insertions(+), 33 deletions(-) diff --git a/tests/migrations/test_writer.py b/tests/migrations/test_writer.py index 248d679a8..03fb22317 100644 --- a/tests/migrations/test_writer.py +++ b/tests/migrations/test_writer.py @@ -77,9 +77,8 @@ def test_writer_format_create_model_basic(tmp_path: Path, monkeypatch) -> None: ] expected = textwrap.dedent( """\ - from tortoise import migrations + from tortoise import fields, migrations from tortoise.migrations import operations as ops - from tortoise import fields class Migration(migrations.Migration): operations = [ @@ -107,9 +106,8 @@ def test_writer_format_rename_and_alter(tmp_path: Path, monkeypatch) -> None: ] expected = textwrap.dedent( """\ - from tortoise import migrations + from tortoise import fields, migrations from tortoise.migrations import operations as ops - from tortoise import fields class Migration(migrations.Migration): operations = [ @@ -156,10 +154,9 @@ def test_writer_format_options_indexes_constraints(tmp_path: Path, monkeypatch) ] expected = textwrap.dedent( """\ - from tortoise import migrations - from tortoise.migrations import operations as ops - from tortoise import fields + from tortoise import fields, migrations from tortoise.indexes import Index, PartialIndex + from tortoise.migrations import operations as ops from tortoise.migrations.constraints import UniqueConstraint class Migration(migrations.Migration): @@ -233,10 +230,9 @@ def test_writer_renders_fk_field(tmp_path: Path, monkeypatch) -> None: ] expected = textwrap.dedent( """\ - from tortoise import migrations - from tortoise.migrations import operations as ops + from tortoise import fields, migrations from tortoise.fields.base import OnDelete - from tortoise import fields + from tortoise.migrations import operations as ops class Migration(migrations.Migration): operations = [ @@ -272,10 +268,9 @@ def test_writer_excludes_fk_source_field(tmp_path: Path, monkeypatch) -> None: ] expected = textwrap.dedent( """\ - from tortoise import migrations - from tortoise.migrations import operations as ops + from tortoise import fields, migrations from tortoise.fields.base import OnDelete - from tortoise import fields + from tortoise.migrations import operations as ops class Migration(migrations.Migration): operations = [ @@ -312,10 +307,9 @@ def test_writer_serializes_on_delete_enum(tmp_path: Path, monkeypatch) -> None: ] expected = textwrap.dedent( """\ - from tortoise import migrations - from tortoise.migrations import operations as ops + from tortoise import fields, migrations from tortoise.fields.base import OnDelete - from tortoise import fields + from tortoise.migrations import operations as ops class Migration(migrations.Migration): operations = [ @@ -341,9 +335,8 @@ def test_writer_skips_missing_db_index(tmp_path: Path, monkeypatch) -> None: ] expected = textwrap.dedent( """\ - from tortoise import migrations + from tortoise import fields, migrations from tortoise.migrations import operations as ops - from tortoise import fields class Migration(migrations.Migration): operations = [ @@ -464,10 +457,9 @@ def test_writer_handles_one_to_one_field(tmp_path: Path, monkeypatch) -> None: ] expected = textwrap.dedent( """\ - from tortoise import migrations - from tortoise.migrations import operations as ops + from tortoise import fields, migrations from tortoise.fields.base import OnDelete - from tortoise import fields + from tortoise.migrations import operations as ops class Migration(migrations.Migration): operations = [ @@ -500,10 +492,9 @@ def test_writer_handles_enum_fields(tmp_path: Path, monkeypatch) -> None: # NOT fields.IntEnumFieldInstance or fields.CharEnumFieldInstance expected = textwrap.dedent( """\ - from tortoise import migrations - from tortoise.migrations import operations as ops from tests.migrations.test_writer import Role, Status - from tortoise import fields + from tortoise import fields, migrations + from tortoise.migrations import operations as ops class Migration(migrations.Migration): operations = [ @@ -533,9 +524,9 @@ def test_writer_format_runpython(tmp_path: Path, monkeypatch) -> None: operations = [RunPython(_runpython_forward, reverse_code=_runpython_reverse, atomic=False)] expected = textwrap.dedent( """\ + from tests.migrations.test_writer import _runpython_forward, _runpython_reverse from tortoise import migrations from tortoise.migrations import operations as ops - from tests.migrations.test_writer import _runpython_forward, _runpython_reverse class Migration(migrations.Migration): operations = [ diff --git a/tortoise/migrations/writer.py b/tortoise/migrations/writer.py index d99332a9f..b134b5112 100644 --- a/tortoise/migrations/writer.py +++ b/tortoise/migrations/writer.py @@ -84,21 +84,28 @@ def add_index_class(self, name: str) -> None: def add_constraint_class(self, name: str) -> None: self.uses_constraints.add(name) - def render(self) -> list[str]: - lines: list[str] = [] + def render(self, *required_imports: str) -> list[str]: + lines: list[str] = [*required_imports] for module in sorted(self.modules): lines.append(f"import {module}") for module, names in sorted(self.imports.items()): lines.append(f"from {module} import {', '.join(sorted(names))}") if self.uses_fields_module: - lines.append("from tortoise import fields") + from_import = "from tortoise import " + for i, line in enumerate(lines): + if line.startswith(from_import): + imported = [i.strip() for i in line[len(from_import) :].split(",")] + lines[i] = from_import + ", ".join(sorted(["fields", *imported])) + break + else: + lines.append("from tortoise import fields") if self.uses_indexes: index_names = ", ".join(sorted(self.uses_indexes)) lines.append(f"from tortoise.indexes import {index_names}") if self.uses_constraints: constraint_names = ", ".join(sorted(self.uses_constraints)) lines.append(f"from tortoise.migrations.constraints import {constraint_names}") - return lines + return sorted(lines) def _resolve_import(value: Any) -> tuple[str, str, bool]: @@ -288,13 +295,11 @@ def as_string(self) -> str: for operation in self.operations: operations.extend(self._format_operation(operation, imports, indent=" " * 8)) - lines: list[str] = [ + required_imports = [ "from tortoise import migrations", "from tortoise.migrations import operations as ops", ] - extra_imports = imports.render() - if extra_imports: - lines.extend(extra_imports) + lines: list[str] = imports.render(*required_imports) lines.extend(["", "class Migration(migrations.Migration):"]) blocks: list[list[str]] = [] if self.dependencies: From 0ff989f0e4603f70aafd2c30cbd12dcc8ffd6ad3 Mon Sep 17 00:00:00 2001 From: Waket Zheng Date: Sat, 16 May 2026 15:34:51 +0800 Subject: [PATCH 2/2] refactor: avoid i shadows --- tortoise/migrations/writer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tortoise/migrations/writer.py b/tortoise/migrations/writer.py index b134b5112..ded70e8df 100644 --- a/tortoise/migrations/writer.py +++ b/tortoise/migrations/writer.py @@ -92,10 +92,10 @@ def render(self, *required_imports: str) -> list[str]: lines.append(f"from {module} import {', '.join(sorted(names))}") if self.uses_fields_module: from_import = "from tortoise import " - for i, line in enumerate(lines): + for idx, line in enumerate(lines): if line.startswith(from_import): imported = [i.strip() for i in line[len(from_import) :].split(",")] - lines[i] = from_import + ", ".join(sorted(["fields", *imported])) + lines[idx] = from_import + ", ".join(sorted(["fields", *imported])) break else: lines.append("from tortoise import fields")