From b5046438d491106622285fb994fd2ece1b4c9b90 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nikola=20Forr=C3=B3?= <nforro@redhat.com>
Date: Thu, 12 Dec 2024 15:32:22 +0100
Subject: [PATCH] Implement `Sections.get_or_create()`
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Nikola Forró <nforro@redhat.com>
---
 specfile/sections.py        | 32 +++++++++++++++++++++++++++++
 tests/unit/test_sections.py | 41 +++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/specfile/sections.py b/specfile/sections.py
index 7b5ae98..37d0ff8 100644
--- a/specfile/sections.py
+++ b/specfile/sections.py
@@ -214,6 +214,38 @@ def find(self, id: str) -> int:
                 return i
         raise ValueError
 
+    def get_or_create(self, id: str) -> Section:
+        """
+        Gets an existing section by ID or creates a new section with the given ID.
+        New section will be appended to the end.
+
+        Args:
+            id: ID of the section (name and options, without the leading '%').
+
+        Returns:
+            Existing or newly created section as an instance of `Section` class.
+        """
+
+        def split_id(id):
+            separator = "\n"
+            tokens = re.split(r"(\s+)", id)
+            if len(tokens) > 2:
+                name = tokens[0]
+                delimiter = tokens[1]
+                options = Options(
+                    Options.tokenize("".join(tokens[2:])),
+                    SECTION_OPTIONS.get(name.lower()),
+                )
+                return name, options, delimiter, separator
+            return tokens[0], None, "", separator
+
+        try:
+            section = self.get(id)
+        except (ValueError, KeyError):
+            section = Section(*split_id(id))
+            self.data.append(section)
+        return section
+
     @classmethod
     def parse(
         cls, lines: List[str], context: Optional["Specfile"] = None
diff --git a/tests/unit/test_sections.py b/tests/unit/test_sections.py
index ccc0736..87bc855 100644
--- a/tests/unit/test_sections.py
+++ b/tests/unit/test_sections.py
@@ -39,6 +39,47 @@ def test_get():
         sections.get("package foo")
 
 
+@pytest.mark.parametrize(
+    "id, existing, name, options, content",
+    [
+        (
+            "package",
+            True,
+            "package",
+            "",
+            ["Name: test", "Version: 0.1", "Release: 1%{?dist}", ""],
+        ),
+        ("prep", True, "prep", "", ["%autosetup", ""]),
+        ("package -n subpkg1", True, "package", "-n subpkg1", [""]),
+        ("package -n subpkg2", False, "package", "-n subpkg2", []),
+    ],
+)
+def test_get_or_create(id, existing, name, options, content):
+    sections = Sections.parse(
+        [
+            "Name: test",
+            "Version: 0.1",
+            "Release: 1%{?dist}",
+            "",
+            "%description",
+            "Test package",
+            "",
+            "%prep",
+            "%autosetup",
+            "",
+            "%package -n subpkg1",
+            "",
+            "%changelog",
+        ]
+    )
+    section = sections.get_or_create(id)
+    assert section.name == name
+    assert str(section.options) == options
+    assert list(section) == content
+    if not existing:
+        assert section == sections[-1]
+
+
 def test_parse():
     sections = Sections.parse(
         [