Move source_to_table to Source class

Secrus · Sep 2, 2024 · 867bc74 · 867bc74
1 parent f843cee
commit 867bc74
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 28 deletions.
diff --git a/src/poetry/config/source.py b/src/poetry/config/source.py
@@ -2,9 +2,12 @@
 
 import dataclasses
 import warnings
+from typing import TYPE_CHECKING
 
 from poetry.repositories.repository_pool import Priority
 
+if TYPE_CHECKING:
+    from tomlkit.items import Table
 
 @dataclasses.dataclass(order=True, eq=True)
 class Source:
@@ -41,3 +44,13 @@ def to_dict(self) -> dict[str, str | bool]:
                 if v
             },
         )
+
+    def to_toml_table(self) -> Table:
+        from tomlkit import nl
+        from tomlkit import table
+
+        source_table: Table = table()
+        for key, value in self.to_dict().items():
+            source_table.add(key, value)
+        source_table.add(nl())
+        return source_table
diff --git a/src/poetry/console/commands/source/add.py b/src/poetry/console/commands/source/add.py
@@ -61,7 +61,6 @@ class SourceAddCommand(Command):
 
     def handle(self) -> int:
         from poetry.factory import Factory
-        from poetry.utils.source import source_to_table
 
         name: str = self.argument("name")
         lower_name = name.lower()
@@ -147,11 +146,11 @@ def handle(self) -> int:
                 source = new_source
                 is_new_source = False
 
-            sources.append(source_to_table(source))
+            sources.append(source.to_toml_table())
 
         if is_new_source:
             self.line(f"Adding source with name <c1>{name}</c1>.")
-            sources.append(source_to_table(new_source))
+            sources.append(new_source.to_toml_table())
         else:
             self.line(f"Source with name <c1>{name}</c1> already exists. Updating.")
 

diff --git a/src/poetry/console/commands/source/remove.py b/src/poetry/console/commands/source/remove.py
@@ -25,8 +25,6 @@ class SourceRemoveCommand(Command):
     ]
 
     def handle(self) -> int:
-        from poetry.utils.source import source_to_table
-
         name = self.argument("name")
         lower_name = name.lower()
 
@@ -38,7 +36,7 @@ def handle(self) -> int:
                 self.line(f"Removing source with name <c1>{source.name}</c1>.")
                 removed = True
                 continue
-            sources.append(source_to_table(source))
+            sources.append(source.to_toml_table())
 
         if not removed:
             self.line_error(

diff --git a/src/poetry/utils/setup_reader.py b/src/poetry/utils/setup_reader.py
diff --git a/src/poetry/utils/source.py b/src/poetry/utils/source.py
diff --git a/tests/utils/test_source.py → tests/config/test_source.py b/tests/utils/test_source.py → tests/config/test_source.py
@@ -8,7 +8,6 @@
 
 from poetry.config.source import Source
 from poetry.repositories.repository_pool import Priority
-from poetry.utils.source import source_to_table
 
 
 @pytest.mark.parametrize(
@@ -36,7 +35,7 @@ def test_source_to_table(source: Source, table_body: dict[str, str | bool]) -> N
     table = Table(Container(), Trivia(), False)
     table._value = table_body  # type: ignore[assignment]
 
-    assert source_to_table(source) == table
+    assert source.to_toml_table() == table
 
 
 def test_source_default_is_primary() -> None: