Commandline plugin - Data subcommand

If your plugin provides custom data types, you might want to provide commandline commands to handle them: Create them from files (example: pseudopotentials), provide export to file formats, visualize them, etc.

With commandline plugins you have the possibility to make you command accessible from the verdi commandline. Your cli plugin will be treated as a subcommand of verdi data.

Exercise: command to export FloatData to file

Plugin structure:

aiida-yourplugin/
   aiida_yourplugin/
      __init__.py
      data/
         __init__.py
         float.py
   setup.py
   setup.json

The file float.py can be taken from the datatype tutorial or replaced by your own custom data type.

File excerpt setup.json:

{
   ...
   "entry_points": {
      "aiida.data": {
         "yourplugin.float = aiida_yourplugin.data.float:FloatData"
      },
      ...
   }
   ...
}

We will assume your plugin provides a FloatData data class. Let’s provide a command that exports it to some file format.

First, we create a new subpackage (this is optional but helps structure our plugin), containing an empty module in which we will work. New plugin structure:

aiida-yourplugin/
   aiida_yourplugin/
      __init__.py
      data/
         __init__.py
         float.py
      cmdline/
         __init__.py
         float_cmd.py  <-- new empty module
   setup.py
   setup.json

Inside that module we will first create an empty command-group (it will do nothing but subcommands can be added to it later). which can be called from the commandline using verdi data yourplugin-float. Command groups are explained in the Click documentation.

File float_cmd.py:

import click  # This we will use in a later step

from aiida.cmdline.commands import data_cmd
from aiida.cmdline.utils.decorators import load_dbenv_if_not_loaded  # Will be used in a later step

@data_cmd.group('yourplugin-float'):
def float_cmd():
   """Commandline interface for working with FloatData"""

This so far does nothing and will not yet be recognized by AiiDA. We will now expose it through an entry point for AiiDA to find. Changes to file setup.json:

{
   ...
   "entry_points": {
      "aiida.data": {
         "yourplugin.float = aiida_yourplugin.data.float:FloatData"
      },
      "aiida.cmdline.data": {                                              <-- NEW
         "yourplugin-float = aiida_yourplugin.cmdline.float_cmd:float_cmd" <-- NEW
      }                                                                    <-- NEW
      ...
   }
   ...
}

Now we only have to reinstall our plugin (pip install -e <path/to/aiida-yourplugin>) and the command should be recognized. We can test it by running:

verdi data yourplugin-float --help

It should print some basic usage information containing the docstring we gave to the float_cmd() function.

The last step is now implementing verdi data yourplugin-float export command that exports our FloatData instance to a file.

Append to file float_cmd.py:

@float_cmd.command()
@click.option('--outfile', '-o', type=click.Path(dir_okay=False), help='write output to this file (by default print to stout).'
@click.argument('pk', type=int)
def export(outfile, pk):
   """Export a FloatData node, identified by PK to plain text format"""
   load_dbenv_if_not_loaded()  # Important to load the dbenv in the last moment
   from aiida.orm import load_node
   float_node = load_node(pk)  # Exercise left to the user: check if it is a FloatData
   file_content = str(float_node.value)
   if outfile:
      with open(outfile, 'w') as out_file_obj:
         out_file_obj.write(file_content)
   else:
      click.echo(file_content)

A subcommand to a group can be defined using the following pattern:

@float_cmd.command()
def export(...):
   ...

Where the subcommand will now automatically have the name of the function. If you want it to have a different name, simply pass it as an argument to the <group>.command('<subcmd name>') decorator.

def export(...):
   """..."""
   load_dbenv_if_not_loaded()  # Important to load the dbenv in the last moment

As is mentioned in the comment, it is important to load the dbenv as late as possible. Particularly it should never be done at import time (on module level) but only inside whichever function requires it. This ensures that command completion does not get slowed down while importing your command.

Last but by no means least, it is important to test our plugin command, this example will use the builtin unittest framework but it is just as well possible to use pytest.

New structure:

aiida-yourplugin/
   aiida_yourplugin/
      __init__.py
      data/
         __init__.py
         float.py
      cmdline/
         __init__.py
         float_cmd.py
         test_float_cmd.py <-- new empty module
   setup.py
   setup.json

Example test in test_float_cmd.py:

import os

from click.testing import CliRunner
from aiida.utils.fixtures import PluginTestCase

from aiida_yourplugin.cmdline.float_cmd import float_cmd

TestFloadCmd(PluginTestCase):
   """Test correctness of the verdi data yourplugin-float export command"""

     BACKEND = os.environ.get('TEST_BACKEND')
     # load the backend to be tested from the environment variable
     # on bash, simply prepend the test command with TEST_BACKEND='django' or TEST_BACKEND='sqlalchemy'
     # or set the TEST_BACKEND in your CI configuration

   def setUp(self):
      from aiida.orm import DataFactory
      self.float_node = DataFactory('yourplugin.float')()
      self.float_node.value = 1.2
      self.runner = CliRunner()

   def test_export(self):
      self.float_node.store()
      result = self.runner.invoke(float_cmd, ['export', str(self.float_node.pk)])
      self.assertEqual(result.output, str(self.float_node.value))

This test can now be run using TEST_BACKEND=django python -m unittest discover from your top level project directory aiida-yourplugin.

As a further exercise, try adding a --format option to choose between plain text and, say json.

Understanding commandline plugins

The discovery of plugins via entry points follows exactly the same mechanisms as all other plugin types.

The possibility of plugging cli commands into each other is a feature of click a python library that greatly simplifies the task. You can find in-depth documentation here: Click 6.0 docs.