CLI: allow folder on disk for aiida-pseudo install family (#80)

Uptil now, the `install family` command only accepted (compressed)
archives of a directory containing the pseudopotential files. Often,
however, a user may already have an unpacked directory on local disk and
it would be ridiculous to force them to archive it first before being
able to use it with the CLI.

The parameter type of the `archive` argument is changed to use the
`PathOrUrl` type of `aiida-core`. This type will treat the value as a
URL and fetch its content if the value is not a valid filepath on the
local file system. This type is not ideal since its error handling is
different from the rest of the `aiida-pseudo` CLI so its a bit
inhomogeneous, but this will be improved later.

Additionally, the `--pseudo-type` option was added. Historically, the
command only exposed the `--family-type` option since in the beginning
of `aiida-pseudo` each family type defined its own hardcoded pseudo type
so it was implied. Since then, families can define their pseudo type per
instance and so now the `--pseudo-type` option has become necessary to
let the user specify which pseudopotential data class should be used.

aiida_pseudo/cli/install.py

@@ -5,6 +5,7 @@
 import pathlib
 import shutil
 import tempfile
+import urllib.request
 
 import click
 
@@ -23,43 +24,55 @@ def cmd_install():
 
 
 @cmd_install.command('family')
-@click.argument('archive', type=types.FileOrUrl(mode='rb'))
+@click.argument('archive', type=types.PathOrUrl(exists=True, file_okay=True))
 @click.argument('label', type=click.STRING)
 @options_core.DESCRIPTION(help='Description for the family.')
 @options.ARCHIVE_FORMAT()
 @options.FAMILY_TYPE()
+@options.PSEUDO_TYPE()
 @options.TRACEBACK()
 @decorators.with_dbenv()
-def cmd_install_family(archive, label, description, archive_format, family_type, traceback):  # pylint: disable=too-many-arguments
-    """Install a standard pseudo potential family from an ARCHIVE on the local file system or from a URL.
+def cmd_install_family(archive, label, description, archive_format, family_type, pseudo_type, traceback):  # pylint: disable=too-many-arguments
+    """Install a standard pseudo potential family from a FOLDER or an ARCHIVE (on the local file system or from a URL).
 
-    The command will attempt to infer the archive format from the filename extension of the ARCHIVE. If this fails, the
-    archive format can be specified explicitly with the archive format option, which will also display which formats
-    are supported.
+    The command will attempt first to recognize the passed ARCHIVE_FOLDER as a folder in the local system. If not,
+    `archive` is assumed to be an archive and the command will attempt to infer the archive format from the
+    filename extension of the ARCHIVE. If this fails, the archive format can be specified explicitly with the archive
+    format option, which will also display which formats are supported.
 
-    By default, the command will create a base `PseudoPotentialFamily`, but the type can be changed with the family
+    By default, the command will create a base `PseudoPotentialFamily`, but the type can be changed with the pseudos
     type option. If the base type is used, the pseudo potential files in the archive *have* to have filenames that
     strictly follow the format `ELEMENT.EXTENSION`, because otherwise the element cannot be determined automatically.
     """
     from .utils import attempt, create_family_from_archive
 
-    # The `archive` is now either a `http.client.HTTPResponse` or a normal filelike object, so we get the original file
-    # name in a different way.
-    try:
-        suffix = os.path.basename(archive.url)
-    except AttributeError:
-        suffix = os.path.basename(archive.name)
-
-    # Write the content of the archive to a temporary file, because `create_family_from_archive` does currently not
-    # accept filelike objects because the underlying `shutil.unpack_archive` does not. Likewise, `unpack_archive` will
-    # attempt to deduce the archive format from the filename extension, so it is important we maintain the original
-    # filename. Of course if this fails, users can specify the archive format explicitly wiht the corresponding option.
-    with tempfile.NamedTemporaryFile(mode='w+b', suffix=suffix) as handle:
-        shutil.copyfileobj(archive, handle)
-        handle.flush()
+    # `archive` is a simple string, containing the name of the folder / file / url location.
 
+    if pathlib.Path(archive).is_dir():
+        try:
+            family = family_type.create_from_folder(archive, label, pseudo_type=pseudo_type)
+        except ValueError as exception:
+            raise OSError(f'failed to parse pseudos from `{archive}`: {exception}') from exception
+    elif pathlib.Path(archive).is_file():
         with attempt('unpacking archive and parsing pseudos... ', include_traceback=traceback):
-            family = create_family_from_archive(family_type, label, pathlib.Path(handle.name), fmt=archive_format)
+            family = create_family_from_archive(
+                family_type, label, pathlib.Path(archive), fmt=archive_format, pseudo_type=pseudo_type
+            )
+    else:
+        # The file of the url must be copied to a local temporary file. Maybe better ways to do it?
+        # The `create_family_from_archive` does currently not accept filelike objects because the underlying
+        # `shutil.unpack_archive` does not. Likewise, `unpack_archive` will attempt to deduce the archive format
+        # from the filename extension, so it is important we maintain the original filename.
+        # Of course if this fails, users can specify the archive format explicitly wiht the corresponding option.
+        with urllib.request.urlopen(archive) as handle:
+            suffix = os.path.basename(handle.url)
+            with tempfile.NamedTemporaryFile(mode='w+b', suffix=suffix) as handle:
+                shutil.copyfileobj(handle, handle)
+                handle.flush()
+                with attempt('unpacking archive and parsing pseudos... ', include_traceback=traceback):
+                    family = create_family_from_archive(
+                        family_type, label, pathlib.Path(handle.name), fmt=archive_format, pseudo_type=pseudo_type
+                    )
 
     family.description = description
     echo.echo_success(f'installed `{label}` containing {family.count()} pseudo potentials')

aiida_pseudo/cli/params/options.py

@@ -5,7 +5,7 @@
 import click
 
 from aiida.cmdline.params.options import OverridableOption
-from .types import PseudoPotentialFamilyTypeParam
+from .types import PseudoPotentialFamilyTypeParam, PseudoPotentialTypeParam
 
 __all__ = (
     'VERSION', 'FUNCTIONAL', 'RELATIVISTIC', 'PROTOCOL', 'PSEUDO_FORMAT', 'STRINGENCY', 'DEFAULT_STRINGENCY',
@@ -69,6 +69,18 @@
     help='Choose the type of pseudo potential family to create.'
 )
 
+PSEUDO_TYPE = OverridableOption(
+    '-P',
+    '--pseudo-type',
+    type=PseudoPotentialTypeParam(),
+    default='pseudo',
+    show_default=True,
+    help=(
+        'Select the pseudopotential type to be used for the family. Should be the entry point name of a '
+        'subclass of `PseudoPotentialData`.'
+    )
+)
+
 ARCHIVE_FORMAT = OverridableOption(
     '-F', '--archive-format', type=click.Choice([fmt[0] for fmt in shutil.get_archive_formats()])
 )

aiida_pseudo/cli/params/types.py

@@ -5,7 +5,45 @@
 
 from aiida.cmdline.params.types import GroupParamType
 
-__all__ = ('PseudoPotentialFamilyTypeParam', 'PseudoPotentialFamilyParam')
+__all__ = ('PseudoPotentialFamilyTypeParam', 'PseudoPotentialFamilyParam', 'PseudoPotentialTypeParam')
+
+
+class PseudoPotentialTypeParam(click.ParamType):
+    """Parameter type for `click` commands to define a subclass of `PseudoPotentialData`."""
+
+    name = 'pseudo_type'
+
+    def convert(self, value, _, __):
+        """Convert the entry point name to the corresponding class.
+
+        :param value: entry point name that should correspond to subclass of `PseudoPotentialData` data plugin
+        :return: the `PseudoPotentialData` subclass
+        :raises: `click.BadParameter` if the entry point cannot be loaded or is not subclass of `PseudoPotentialData`
+        """
+        from aiida.common import exceptions
+        from aiida.plugins import DataFactory
+        from aiida_pseudo.data.pseudo import PseudoPotentialData
+
+        try:
+            pseudo_type = DataFactory(value)
+        except exceptions.EntryPointError as exception:
+            raise click.BadParameter(f'`{value}` is not an existing data plugin.') from exception
+
+        if not issubclass(pseudo_type, PseudoPotentialData):
+            raise click.BadParameter(f'`{value}` entry point is not a subclass of `PseudoPotentialData`.')
+
+        PseudoPotentialData.entry_point = value
+
+        return pseudo_type
+
+    def complete(self, _, incomplete):
+        """Return possible completions based on an incomplete value.
+
+        :returns: list of tuples of valid entry points (matching incomplete) and a description
+        """
+        from aiida.plugins.entry_point import get_entry_point_names
+        entry_points = get_entry_point_names('aiida.data')
+        return [(ep, '') for ep in entry_points if (ep.startswith('pseudo') and ep.startswith(incomplete))]
 
 
 class PseudoPotentialFamilyParam(GroupParamType):

tests/cli/test_install.py

@@ -10,6 +10,7 @@
 
 from aiida_pseudo.cli import install
 from aiida_pseudo.cli import cmd_install_family, cmd_install_sssp, cmd_install_pseudo_dojo
+from aiida_pseudo.data.pseudo.upf import UpfData
 from aiida_pseudo.groups.family import PseudoPotentialFamily
 from aiida_pseudo.groups.family.pseudo_dojo import PseudoDojoFamily, PseudoDojoConfiguration
 from aiida_pseudo.groups.family.sssp import SsspFamily, SsspConfiguration
@@ -129,6 +130,24 @@ def test_install_family(run_cli_command, get_pseudo_archive):
     assert len(family.pseudos) != 0
 
 
+@pytest.mark.usefixtures('clear_db')
+def test_install_family_folder(run_cli_command, filepath_pseudos):
+    """Test ``aiida-pseudo install family` from folder`."""
+    label = 'family_test'
+    description = 'description'
+    dirpath = filepath_pseudos()
+    options = ['-D', description, dirpath, label]
+
+    result = run_cli_command(cmd_install_family, options)
+    assert f'installed `{label}`' in result.output
+    assert PseudoPotentialFamily.objects.count() == 1
+
+    family = PseudoPotentialFamily.objects.get(label=label)
+    assert family.__class__ is PseudoPotentialFamily
+    assert family.description == description
+    assert len(family.pseudos) != 0
+
+
 @pytest.mark.usefixtures('clear_db')
 def test_install_family_url(run_cli_command):
     """Test ``aiida-pseudo install family`` when installing from a URL.
@@ -139,14 +158,15 @@ def test_install_family_url(run_cli_command):
     label = 'SSSP/1.0/PBE/efficiency'
     description = 'description'
     filepath_archive = 'https://legacy-archive.materialscloud.org/file/2018.0001/v4/SSSP_1.0_PBE_efficiency.tar.gz'
-    options = ['-D', description, filepath_archive, label, '-T', 'pseudo.family.sssp']
+    options = ['-D', description, filepath_archive, label, '-P', 'pseudo.upf']
 
     result = run_cli_command(cmd_install_family, options)
     assert f'installed `{label}`' in result.output
-    assert SsspFamily.objects.count() == 1
+    assert PseudoPotentialFamily.objects.count() == 1
 
-    family = SsspFamily.objects.get(label=label)
-    assert family.__class__ is SsspFamily
+    family = PseudoPotentialFamily.objects.get(label=label)
+    assert isinstance(family.pseudos['Si'], UpfData)
+    assert family.__class__ is PseudoPotentialFamily
     assert family.description == description
     assert len(family.pseudos) != 0
 
@@ -252,8 +272,6 @@ def test_install_sssp_download_only_exists(run_monkeypatched_install_sssp, get_p
 
     The files should be downloadable even if the corresponding configuration is already installed.
     """
-    from aiida_pseudo.data.pseudo.upf import UpfData
-
     version = '1.1'
     functional = 'PBEsol'
     protocol = 'precision'
@@ -285,8 +303,6 @@ def test_install_pseudo_dojo_download_only_exists(run_monkeypatched_install_pseu
 
     The files should be downloadable even if the corresponding configuration is already installed.
     """
-    from aiida_pseudo.data.pseudo.upf import UpfData
-
     version = '1.0'
     functional = 'LDA'
     relativistic = 'SR'