From 1d4da7f41a5a0403409ab6eca6ce656674000cc0 Mon Sep 17 00:00:00 2001
From: Matt Culler <matt.culler@canonical.com>
Date: Thu, 29 Aug 2024 15:34:09 -0400
Subject: [PATCH] feat(partitions): custom error for nonexistent partition
 (#826)

* feat(partitions): custom error for nonexistent partition

* chore: deduplicate checks and make type checker happy

* fix: had been making calls incorrectly

* chore: update wordlist

* Update craft_parts/dirs.py

Co-authored-by: Alex Lowe <alex.lowe@canonical.com>

* docs: clarify variable's purpose

* refactor: autoformat

---------

Co-authored-by: Alex Lowe <alex.lowe@canonical.com>
---
 craft_parts/dirs.py                           | 24 +++++++++++++----
 craft_parts/errors.py                         | 26 +++++++++++++++++--
 .../craft-parts/craft-parts.wordlist.txt      |  1 +
 3 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/craft_parts/dirs.py b/craft_parts/dirs.py
index aa2ed51c9..1d28ca1ce 100644
--- a/craft_parts/dirs.py
+++ b/craft_parts/dirs.py
@@ -20,6 +20,7 @@
 from pathlib import Path
 from types import MappingProxyType
 
+from craft_parts.errors import PartitionNotFound, PartitionUsageError
 from craft_parts.utils import partition_utils
 
 
@@ -74,14 +75,27 @@ def __init__(
             )
         )
 
+    def _validate_requested_partition(
+        self, dir_name: str, partition: str | None = None
+    ) -> None:
+        """Ensure the requested partition is valid."""
+        if self._partitions:
+            if not partition:
+                raise PartitionUsageError(
+                    error_list=[
+                        f"Partitions are enabled, you must specify which partition's {dir_name!r} you want."
+                    ],
+                    partitions=self._partitions,
+                )
+            if partition not in self._partitions:
+                raise PartitionNotFound(partition, self._partitions)
+
     def get_stage_dir(self, partition: str | None = None) -> Path:
         """Get the stage directory for the given partition."""
-        if self._partitions and partition not in self._partitions:
-            raise ValueError(f"Unknown partition {partition}")
+        self._validate_requested_partition("stage_dir", partition)
         return self.stage_dirs[partition]
 
     def get_prime_dir(self, partition: str | None = None) -> Path:
-        """Get the stage directory for the given partition."""
-        if self._partitions and partition not in self._partitions:
-            raise ValueError(f"Unknown partition {partition}")
+        """Get the prime directory for the given partition."""
+        self._validate_requested_partition("prime_dir", partition)
         return self.prime_dirs[partition]
diff --git a/craft_parts/errors.py b/craft_parts/errors.py
index 09fa43df0..39008b1a4 100644
--- a/craft_parts/errors.py
+++ b/craft_parts/errors.py
@@ -652,17 +652,22 @@ class PartitionUsageError(PartitionError):
     """Error for a list of invalid partition usages.
 
     :param error_list: Iterable of strings describing the invalid usages.
+    :param partitions: Iterable of the names of valid partitions.
+    :param brief: Override brief message.
     """
 
     def __init__(
-        self, error_list: Iterable[str], partitions: Iterable[str] | None
+        self,
+        error_list: Iterable[str],
+        partitions: Iterable[str] | None,
+        brief: str | None = None,
     ) -> None:
         valid_partitions = (
             f"\nValid partitions: {', '.join(partitions)}" if partitions else ""
         )
 
         super().__init__(
-            brief="Invalid usage of partitions",
+            brief=brief or "Invalid usage of partitions",
             details="\n".join(error_list) + valid_partitions,
             resolution="Correct the invalid partition name(s) and try again.",
         )
@@ -688,3 +693,20 @@ def __init__(self, warning_list: Iterable[str]) -> None:
             ),
         )
         Warning.__init__(self)
+
+
+class PartitionNotFound(PartitionUsageError):
+    """A partition has been specified that does not exist.
+
+    :param partition_name: The name of the partition that does not exist.
+    :param partitions: Iterable of the names of valid partitions.
+    """
+
+    def __init__(self, partition_name: str, partitions: Iterable[str]) -> None:
+        # Allow callers catching this exception easy access to the partition name
+        self.partition_name = partition_name
+        super().__init__(
+            brief=f"Requested partition does not exist: {partition_name!r}",
+            partitions=partitions,
+            error_list=[],
+        )
diff --git a/docs/common/craft-parts/craft-parts.wordlist.txt b/docs/common/craft-parts/craft-parts.wordlist.txt
index a9aace323..7add304af 100644
--- a/docs/common/craft-parts/craft-parts.wordlist.txt
+++ b/docs/common/craft-parts/craft-parts.wordlist.txt
@@ -175,6 +175,7 @@ PartInfo
 PartSpec
 PartSpecificationError
 PartitionError
+PartitionNotFound
 PartitionPathPair
 PartitionUsageError
 PartitionUsageWarning