add tests and clean up function

craft_parts/utils/partition_utils.py

@@ -17,17 +17,23 @@
 """Partition helpers."""
 
 import re
-from pathlib import Path
-from typing import Union
+from pathlib import PurePath, PurePosixPath
+from typing import TypeVar, Union
 
 from craft_parts.features import Features
 
+FlexiblePath = TypeVar("FlexiblePath", bound=Union[PurePath, str])
 
-# XXX: for context, you can see where this function in the prototype:
-# https://github.com/mr-cal/craft-parts/commit/e2d15f1bb1da3207463ad3015b68509537b007c3
 
+HAS_PARTITION_REGEX = re.compile(r"^\([a-z]+\)(/.*)?$")
 
-def get_partition_compatible_filepath(filepath: Union[Path, str]) -> Union[Path, str]:
+
+def _has_partition(path: FlexiblePath) -> bool:
+    """Check whether a path has an explicit partition."""
+    return bool(HAS_PARTITION_REGEX.match(str(path)))
+
+
+def get_partition_compatible_filepath(filepath: FlexiblePath) -> FlexiblePath:
     """Get a filepath compatible with the partitions feature.
 
     If the filepath begins with a partition, then the parentheses are stripped from the
@@ -43,25 +49,22 @@ def get_partition_compatible_filepath(filepath: Union[Path, str]) -> Union[Path,
     if not Features().enable_partitions:
         return filepath
 
-    # XXX: I think this is the correct thing to do. The default globs should not be
-    # modified (which comes from craft_parts/parts.py lines 55-58)
-    if filepath == "*":
+    str_path = str(filepath)
+
+    # Anything that starts with a glob should be assumed to remain that same glob.
+    if str_path.startswith("*"):
         return filepath
 
-    # XXX: I know there is at least one problem with this regex: `(default)` and
-    # `(default)/` are treated differently
-    match = re.match("^\\((?P<partition>[a-z]+)\\)/(?P<filepath>.*)", str(filepath))
-    if match:
-        partition = match.group("partition")
-        everything_else = match.group("filepath")
+    if _has_partition(str_path):
+        if "/" not in str_path:
+            partition = str_path.strip("()")
+            inner_path = ""
+        else:
+            partition, inner_path = str_path.split("/", maxsplit=1)
+            partition = partition.strip("()")
     else:
         partition = "default"
-        everything_else = filepath
-
-    new_filepath = Path(partition, everything_else)
+        inner_path = str_path
 
-    # XXX: craft-parts calls this function in different places. Sometimes the filepath
-    # is a string and something it is a Path object. I decided to make this function
-    # handle both scenarios but I wonder if this approach will create
-    # type-checking issues in the future.
-    return str(new_filepath) if isinstance(filepath, str) else new_filepath
+    new_filepath = PurePosixPath(partition, inner_path)
+    return filepath.__class__(new_filepath)

tests/unit/utils/test_partition_utils.py

@@ -0,0 +1,107 @@
+# -*- Mode:Python; indent-tabs-mode:nil; tab-width:4 -*-
+#
+# Copyright 2023 Canonical Ltd.
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License version 3 as published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""Unit tests for partition utilities."""
+from pathlib import Path, PurePosixPath
+
+import pytest
+from utils.partition_utils import _has_partition, get_partition_compatible_filepath
+
+PATH_CLASSES = [Path, PurePosixPath, str]
+
+NON_PARTITION_PATHS = [
+    "/absolute/path",
+    "relative/path",
+    "",
+]
+
+PARTITION_PATHS = [
+    "(default)",
+    "(default)/",
+    "(default)/path",
+    "(partition)/path",
+]
+
+PARTITION_EXPECTED_PATHS = [
+    "default",
+    "default",
+    "default/path",
+    "partition/path",
+]
+
+# Prevent us from adding nonmatching paths for tests below.
+assert len(PARTITION_PATHS) == len(
+    PARTITION_EXPECTED_PATHS
+), "Expected partition paths and input partition paths need to match 1:1"
+
+
+@pytest.mark.parametrize(
+    ("full_path", "expected"),
+    [
+        ("some/path", False),
+        ("(default)", True),
+        ("(default)/", True),
+        ("(part)/some/path", True),
+        ("(nota)partition", False),
+    ],
+)
+def test_has_partition(full_path, expected):
+    """Test that the partition regex has the expected results."""
+    assert _has_partition(full_path) == expected
+
+
+@pytest.mark.parametrize("path", NON_PARTITION_PATHS + PARTITION_PATHS)
+@pytest.mark.parametrize("cls", PATH_CLASSES)
+def test_get_partition_compatible_filepath_disabled_passthrough(path, cls):
+    """Test that when partitions are disabled this is a no-op."""
+    actual = get_partition_compatible_filepath(cls(path))
+
+    assert actual == cls(path)
+    assert isinstance(actual, cls)
+
+
+@pytest.mark.parametrize("path", ["*", "**", "*/path"])
+@pytest.mark.parametrize("cls", PATH_CLASSES)
+@pytest.mark.usefixtures("enable_partitions_feature")
+def test_get_partition_compatible_filepath_glob_start(path, cls):
+    expected = cls(path)
+    actual = get_partition_compatible_filepath(expected)
+
+    assert actual == expected
+
+
+@pytest.mark.parametrize("path", NON_PARTITION_PATHS)
+@pytest.mark.parametrize("cls", PATH_CLASSES)
+@pytest.mark.usefixtures("enable_partitions_feature")
+def test_get_partition_compatible_filepath_non_partition(path, cls):
+    """Non-partitioned paths get a default partition."""
+    actual = get_partition_compatible_filepath(cls(path))
+
+    assert actual == cls(PurePosixPath("default", path))
+    assert isinstance(actual, cls)
+
+
+@pytest.mark.parametrize(
+    ("path", "expected"),
+    zip(PARTITION_PATHS, PARTITION_EXPECTED_PATHS),
+)
+@pytest.mark.parametrize("cls", PATH_CLASSES)
+@pytest.mark.usefixtures("enable_partitions_feature")
+def test_get_partition_compatible_filepath_partition(path, expected, cls):
+    """Non-partitioned paths match their given partition."""
+    actual = get_partition_compatible_filepath(cls(path))
+
+    assert actual == cls(expected)
+    assert isinstance(actual, cls)