From edd2c3c83bf8b66547099872056f1b07e9ba7091 Mon Sep 17 00:00:00 2001
From: Martin Lehmann <martin.lehmann@deutschebahn.com>
Date: Fri, 30 Jun 2023 11:12:40 +0200
Subject: [PATCH] feat(loader): Derive entrypoint from filehandler's listing

This commit makes the MelodyLoader use the new FileHandler capability to
enumerate files in order to automatically derive the `entrypoint` in
more cases. It will now look into the filehandler's root directory, and
if there is exactly one `.aird` file, that becomes the new entrypoint.

If there is more than one `.aird` file in the root directory, an error
will be raised, and the same is true if the root directory does not
directly contain any `.aird` files. In these cases, it is still
necessary to manually specify the `entrypoint` at load time.
---
 capellambse/loader/core.py  | 61 ++++++++++++++++++++++++-------------
 tests/test_model_loading.py |  2 +-
 2 files changed, 40 insertions(+), 23 deletions(-)

diff --git a/capellambse/loader/core.py b/capellambse/loader/core.py
index b77baeb69..3a329ffd0 100644
--- a/capellambse/loader/core.py
+++ b/capellambse/loader/core.py
@@ -30,7 +30,6 @@
 import capellambse
 import capellambse._namespaces as _n
 from capellambse import filehandler, helpers
-from capellambse.filehandler import local
 from capellambse.loader import exs
 from capellambse.loader.modelinfo import ModelInfo
 
@@ -86,6 +85,41 @@
 METADATA_TAG = f"{{{_n.NAMESPACES['metadata']}}}Metadata"
 
 
+def _derive_entrypoint(
+    path: str | os.PathLike | filehandler.FileHandler,
+    entrypoint: str | pathlib.PurePosixPath | None = None,
+) -> tuple[filehandler.FileHandler, pathlib.PurePosixPath]:
+    if entrypoint:
+        if not isinstance(path, filehandler.FileHandler):
+            path = filehandler.get_filehandler(path)
+        entrypoint = helpers.normalize_pure_path(entrypoint)
+        return path, entrypoint
+
+    if not isinstance(path, filehandler.FileHandler):
+        path = os.fspath(path)
+        protocol, nested_path = filehandler.split_protocol(path)
+        if protocol == "file":
+            assert isinstance(nested_path, pathlib.Path)
+            if nested_path.suffix == ".aird":
+                entrypoint = pathlib.PurePosixPath(nested_path.name)
+                path = nested_path.parent
+                return filehandler.get_filehandler(path), entrypoint
+            elif nested_path.is_file():
+                raise ValueError(
+                    f"Invalid entrypoint: Not an .aird file: {nested_path}"
+                )
+        path = filehandler.get_filehandler(path)
+
+    aird_files = [i for i in path.iterdir() if i.name.endswith(".aird")]
+    if not aird_files:
+        raise ValueError("No .aird file found, specify entrypoint")
+    if len(aird_files) > 1:
+        raise ValueError("Multiple .aird files found, specify entrypoint")
+    entrypoint = pathlib.PurePosixPath(aird_files[0])
+
+    return path, entrypoint
+
+
 def _verify_extension(filename: pathlib.PurePosixPath) -> None:
     """Check whether ``filename`` has a valid extension."""
     file = pathlib.PurePosixPath(filename)
@@ -382,10 +416,10 @@ def __init__(
             "ignore_duplicate_uuids_and_void_all_warranties", False
         )
 
-        if isinstance(path, filehandler.FileHandler):
-            handler = path
-        else:
-            handler = filehandler.get_filehandler(path, **kwargs)
+        handler, self.entrypoint = _derive_entrypoint(path, entrypoint)
+        if self.entrypoint.suffix != ".aird":
+            raise ValueError("Invalid entrypoint, specify the ``.aird`` file")
+
         self.resources = ResourceLocationManager({"\0": handler})
         for resname, reshdl in (resources or {}).items():
             if not resname:
@@ -399,9 +433,6 @@ def __init__(
                 self.resources[resname] = filehandler.get_filehandler(**reshdl)
             else:
                 self.resources[resname] = reshdl
-        self.entrypoint = self.__derive_entrypoint(entrypoint)
-        if self.entrypoint.suffix != ".aird":
-            raise ValueError("Invalid entrypoint, specify the ``.aird`` file")
 
         self.trees: dict[pathlib.PurePosixPath, ModelFile] = {}
         self.__load_referenced_files(
@@ -432,20 +463,6 @@ def check_duplicate_uuids(self):
                 " - check the 'resources' for duplicate models"
             )
 
-    def __derive_entrypoint(
-        self, entrypoint: str | pathlib.PurePosixPath | None
-    ) -> pathlib.PurePosixPath:
-        if entrypoint:
-            return helpers.normalize_pure_path(entrypoint)
-
-        if isinstance(self.filehandler, local.LocalFileHandler):
-            basedir = self.filehandler.path
-            assert isinstance(basedir, pathlib.Path)
-            self.filehandler.path = basedir.parent
-            return helpers.normalize_pure_path(basedir.name)
-
-        raise ValueError("This type of file handler needs an ``entrypoint``")
-
     def __load_referenced_files(
         self, resource_path: pathlib.PurePosixPath
     ) -> None:
diff --git a/tests/test_model_loading.py b/tests/test_model_loading.py
index 0ffa50e68..7ddedbffe 100644
--- a/tests/test_model_loading.py
+++ b/tests/test_model_loading.py
@@ -45,7 +45,7 @@ def test_model_loading_via_LocalFileHandler(path: str | pathlib.Path):
     capellambse.MelodyModel(path)
 
 
-@pytest.mark.parametrize("suffix", [".capella", ".melodymodel"])
+@pytest.mark.parametrize("suffix", [".afm", ".capella"])
 def test_model_loading_with_invalid_entrypoint_fails(suffix: str):
     with pytest.raises(ValueError, match="(?i)invalid entrypoint"):
         capellambse.MelodyModel(TEST_MODEL_5_0.with_suffix(suffix))