abstract search for files in a data directory

Instead of passing down just a str|None, create a class that
looks for files both with pkg_resources and in the user
data directory.

Signed-off-by: Paolo Bonzini <pbonzini@redhat.com>

3 files changed, 59 insertions, 20 deletions

diff --git a/lcitool/commandline.py b/lcitool/commandline.py
index f03aaeb..4555b15 100644
--- a/lcitool/commandline.py
+++ b/lcitool/commandline.py
@@ -11,10 +11,22 @@ import argparse
 from pathlib import Path
 
 from lcitool.application import Application
+from lcitool.util import DataDir
+
 
 log = logging.getLogger(__name__)
 
 
+class DataDirAction(argparse.Action):
+    def __init__(self, option_strings, dest, default=DataDir(), nargs=None, **kwargs):
+        if nargs is not None:
+            raise ValueError("nargs not allowed")
+        super().__init__(option_strings, dest, default=default, nargs=1, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        setattr(namespace, self.dest, DataDir(values[0]))
+
+
 class CommandLine:
 
     def __init__(self):
@@ -204,6 +216,7 @@ class CommandLine:
         )
         self._parser.add_argument(
             "-d", "--data-dir",
+            action=DataDirAction,
             help="extra directory for loading data files from")
 
         subparsers = self._parser.add_subparsers(metavar="ACTION",
diff --git a/lcitool/projects.py b/lcitool/projects.py
index c99f71b..49b4cca 100644
--- a/lcitool/projects.py
+++ b/lcitool/projects.py
@@ -7,11 +7,9 @@
 import logging
 import yaml
 
-from pathlib import Path
-from pkg_resources import resource_filename
-
 from lcitool import util, LcitoolError
 from lcitool.packages import PyPIPackage, CPANPackage
+from lcitool.util import DataDir
 
 log = logging.getLogger(__name__)
 
@@ -52,35 +50,27 @@ class Projects:
             self._load_internal()
         return self._internal
 
-    def __init__(self, data_dir=None):
+    def __init__(self, data_dir=DataDir()):
         self._data_dir = data_dir
         self._public = None
         self._internal = None
 
-    def _load_projects_from_path(self, path):
+    def _load_projects_from_files(self, files):
         projects = {}
 
-        for item in path.iterdir():
-            if not item.is_file() or item.suffix != ".yml":
-                continue
-
-            projects[item.stem] = Project(self, item.stem, item)
+        for item in files:
+            if item.stem not in projects:
+                projects[item.stem] = Project(self, item.stem, item)
 
         return projects
 
     def _load_public(self):
-        source = Path(resource_filename(__name__, "facts/projects"))
-        projects = self._load_projects_from_path(source)
-
-        if self._data_dir is not None:
-            source = Path(self._data_dir).joinpath("projects")
-            projects.update(self._load_projects_from_path(source))
-
-        self._public = projects
+        files = self._data_dir.list_files("facts/projects", ".yml")
+        self._public = self._load_projects_from_files(files)
 
     def _load_internal(self):
-        source = Path(resource_filename(__name__, "facts/projects/internal"))
-        self._internal = self._load_projects_from_path(source)
+        files = self._data_dir.list_files("facts/projects/internal", ".yml", internal=True)
+        self._internal = self._load_projects_from_files(files)
 
     def expand_names(self, pattern):
         try:
diff --git a/lcitool/util.py b/lcitool/util.py
index 2de718e..4fb65a1 100644
--- a/lcitool/util.py
+++ b/lcitool/util.py
@@ -13,6 +13,7 @@ import tempfile
 import textwrap
 
 from pathlib import Path
+from pkg_resources import resource_filename
 
 _tempdir = None
 
@@ -217,6 +218,41 @@ def merge_dict(source, dest):
             merge_dict(source[key], dest[key])
 
 
+class DataDir:
+    """A class that looks for files both under the lcitool sources and in
+       an externally specified data directory.  Used to implement the
+       -d option."""
+
+    def __init__(self, extra_data_dir=None):
+        self._extra_data_dir = extra_data_dir
+
+    def __repr__(self):
+        return f'DataDir({str(self._extra_data_dir)})'
+
+    def _search(self, resource_path, *names, internal=False):
+        if not internal and self._extra_data_dir:
+            # The first part of the path is used to keep data files out of
+            # the source directory, for example "facts" or "etc".  Remove it
+            # when using an external data directory.
+            if "/" in resource_path:
+                user_path = resource_path[resource_path.index("/") + 1:]
+            else:
+                user_path = ""
+            p = Path(self._extra_data_dir, user_path, *names)
+            if p.exists():
+                yield p
+
+        p = Path(resource_filename(__name__, resource_path), *names)
+        if p.exists():
+            yield p
+
+    def list_files(self, resource_path, suffix=None, internal=False):
+        for p in self._search(resource_path, internal=internal):
+            for file in p.iterdir():
+                if file.is_file() and (suffix is None or file.suffix == suffix):
+                    yield file
+
+
 def validate_cross_platform(cross_arch, osname):
     if osname not in ["Debian", "Fedora"]:
         raise ValueError(f"Cannot cross compile on {osname}")