Fix autocompletion for LinkManager and AttributeManager

aiida/common/exceptions.py

@@ -10,14 +10,14 @@
 """Module that define the exceptions that are thrown by AiiDA's internal code."""
 
 __all__ = (
-    'AiidaException', 'NotExistent', 'MultipleObjectsError', 'RemoteOperationError', 'ContentNotExistent',
-    'FailedError', 'StoringNotAllowed', 'ModificationNotAllowed', 'IntegrityError', 'UniquenessError',
-    'EntryPointError', 'MissingEntryPointError', 'MultipleEntryPointError', 'LoadingEntryPointError',
-    'InvalidEntryPointTypeError', 'InvalidOperation', 'ParsingError', 'InternalError', 'PluginInternalError',
-    'ValidationError', 'ConfigurationError', 'ProfileConfigurationError', 'MissingConfigurationError',
-    'ConfigurationVersionError', 'IncompatibleDatabaseSchema', 'DbContentError', 'InputValidationError',
-    'FeatureNotAvailable', 'FeatureDisabled', 'LicensingException', 'TestsNotAllowedError', 'UnsupportedSpeciesError',
-    'TransportTaskException', 'OutputParsingError'
+    'AiidaException', 'NotExistent', 'NotExistentAttributeError', 'NotExistentKeyError', 'MultipleObjectsError',
+    'RemoteOperationError', 'ContentNotExistent', 'FailedError', 'StoringNotAllowed', 'ModificationNotAllowed',
+    'IntegrityError', 'UniquenessError', 'EntryPointError', 'MissingEntryPointError', 'MultipleEntryPointError',
+    'LoadingEntryPointError', 'InvalidEntryPointTypeError', 'InvalidOperation', 'ParsingError', 'InternalError',
+    'PluginInternalError', 'ValidationError', 'ConfigurationError', 'ProfileConfigurationError',
+    'MissingConfigurationError', 'ConfigurationVersionError', 'IncompatibleDatabaseSchema', 'DbContentError',
+    'InputValidationError', 'FeatureNotAvailable', 'FeatureDisabled', 'LicensingException', 'TestsNotAllowedError',
+    'UnsupportedSpeciesError', 'TransportTaskException', 'OutputParsingError'
 )
 
 
@@ -36,6 +36,18 @@ class NotExistent(AiidaException):
     """
 
 
+class NotExistentAttributeError(AttributeError, NotExistent):
+    """
+    Raised when the required entity does not exist, when fetched as an attribute.
+    """
+
+
+class NotExistentKeyError(KeyError, NotExistent):
+    """
+    Raised when the required entity does not exist, when fetched as a dictionary key.
+    """
+
+
 class MultipleObjectsError(AiidaException):
     """
     Raised when more than one entity is found in the DB, but only one was

aiida/orm/utils/managers.py

@@ -14,6 +14,7 @@
 """
 
 from aiida.common.links import LinkType
+from aiida.common.exceptions import NotExistent, NotExistentAttributeError, NotExistentKeyError
 
 __all__ = ('NodeLinksManager', 'AttributeManager')
 
@@ -80,8 +81,15 @@ def __getattr__(self, name):
         """
         try:
             return self._get_node_by_link_label(label=name)
-        except KeyError:
-            raise AttributeError("Node '{}' does not have an input with link '{}'".format(self._node.pk, name))
+        except NotExistent:
+            # Note: in order for TAB-completion to work, we need to raise an
+            # exception that also inherits from AttributeError, so that
+            # `getattr(node.inputs, 'some_label', some_default)` returns
+            # `some_default`. Otherwise, the exception is not caught by
+            # `getattr` and is just propagated, instead of returning the default.
+            raise NotExistentAttributeError(
+                "Node '{}' does not have an input with link '{}'".format(self._node.pk, name)
+            )
 
     def __getitem__(self, name):
         """
@@ -91,8 +99,12 @@ def __getitem__(self, name):
         """
         try:
             return self._get_node_by_link_label(label=name)
-        except KeyError:
-            raise KeyError("Node '{}' does not have an input with link '{}'".format(self._node.pk, name))
+        except NotExistent:
+            # Note: in order for this class to behave as a dictionary, we raise
+            # an exception that also inherits from KeyError - in this way, users
+            # can use the standard construct `try/except KeyError` and this will
+            # behave like a standard dictionary.
+            raise NotExistentKeyError("Node '{}' does not have an input with link '{}'".format(self._node.pk, name))
 
     def __str__(self):
         """Return a string representation of the manager"""
@@ -125,7 +137,7 @@ def __dir__(self):
         """
         Allow to list the keys of the dictionary
         """
-        return sorted(self._node.attributes_items())
+        return sorted(self._node.attributes_keys())
 
     def __iter__(self):
         """

aiida/parsers/plugins/arithmetic/add.py

@@ -7,18 +7,18 @@
 # For further information on the license, see the LICENSE.txt file        #
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
-from aiida.common import exceptions
 from aiida.orm import Int
 from aiida.parsers.parser import Parser
 
 
 class ArithmeticAddParser(Parser):
+    """Parser for an ArithmeticAdd calculation."""
 
-    def parse(self, **kwargs):
+    def parse(self, **kwargs):  # pylint: disable=inconsistent-return-statements
         """Parse the contents of the output files retrieved in the `FolderData`."""
         try:
             output_folder = self.retrieved
-        except exceptions.NotExistent:
+        except AttributeError:
             return self.exit_codes.ERROR_NO_RETRIEVED_FOLDER
 
         try:
@@ -32,7 +32,7 @@ def parse(self, **kwargs):
 
         try:
             allow_negative = self.node.inputs.settings.get_attribute('allow_negative', True)
-        except exceptions.NotExistent:
+        except AttributeError:
             allow_negative = True
 
         if not allow_negative and result < 0:

tests/orm/node/test_node.py

@@ -706,29 +706,29 @@ def test_tab_completable_properties(self):
         # .inputs for calculations
         self.assertEqual(calc1.inputs.input_value.pk, input1.pk)
         self.assertEqual(calc2.inputs.input_value.pk, input2.pk)
-        with self.assertRaises(exceptions.NotExistent):
+        with self.assertRaises(AttributeError):
             _ = calc1.inputs.some_label
 
         # .inputs for workflows
         self.assertEqual(top_workflow.inputs.a.pk, input1.pk)
         self.assertEqual(top_workflow.inputs.b.pk, input2.pk)
         self.assertEqual(workflow.inputs.a.pk, input1.pk)
         self.assertEqual(workflow.inputs.b.pk, input2.pk)
-        with self.assertRaises(exceptions.NotExistent):
+        with self.assertRaises(AttributeError):
             _ = workflow.inputs.some_label
 
         # .outputs for calculations
         self.assertEqual(calc1.outputs.result.pk, output1.pk)
         self.assertEqual(calc2.outputs.result.pk, output2.pk)
-        with self.assertRaises(exceptions.NotExistent):
+        with self.assertRaises(AttributeError):
             _ = calc1.outputs.some_label
 
         # .outputs for workflows
         self.assertEqual(top_workflow.outputs.result_a.pk, output1.pk)
         self.assertEqual(top_workflow.outputs.result_b.pk, output2.pk)
         self.assertEqual(workflow.outputs.result_a.pk, output1.pk)
         self.assertEqual(workflow.outputs.result_b.pk, output2.pk)
-        with self.assertRaises(exceptions.NotExistent):
+        with self.assertRaises(AttributeError):
             _ = workflow.outputs.some_label
 
 

tests/orm/utils/test_managers.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""Tests for the various node managers (.inputs, .outputs, .dict, ...)."""
+# pylint: disable=unused-argument
+
+import pytest
+
+from aiida import orm
+from aiida.common.exceptions import NotExistent, NotExistentAttributeError, NotExistentKeyError
+from aiida.common import LinkType
+
+
+def test_dot_dict_manager(clear_database_before_test):
+    """Verify that the Dict.dict manager behaves as intended."""
+    dict_content = {'a': True, 'b': 1, 'c': 'Some string'}
+    dict_node = orm.Dict(dict=dict_content)
+
+    # Check that dir() return all keys and nothing else, important
+    # for tab competion
+    assert len(dir(dict_node.dict)) == len(dict_content)
+    assert set(dir(dict_node.dict)) == set(dict_content)
+    # Check that it works also as an iterator
+    assert len(list(dict_node.dict)) == len(dict_content)
+    assert set(dict_node.dict) == set(dict_content)
+
+    for key, val in dict_content.items():
+        # dict_node.dict.a == True, ...
+        assert getattr(dict_node.dict, key) == val
+        # dict_node.dict['a'] == True, ...
+        assert dict_node.dict[key] == val
+
+    # I check the attribute fetching directly
+    assert dict_node.dict.b == 1
+
+    # Must raise a AttributeError, otherwise tab competion will not work
+    with pytest.raises(AttributeError):
+        getattr(dict_node.dict, 'NotExistentKey')
+
+    # Must raise a KeyError
+    with pytest.raises(KeyError):
+        _ = dict_node.dict['NotExistentKey']
+
+
+def test_link_manager(clear_database_before_test):
+    """Test the LinkManager via .inputs and .outputs from a ProcessNode."""
+    # I first create a calculation with two inputs and two outputs
+
+    # Create inputs
+    inp1 = orm.Data()
+    inp1.store()
+    inp2 = orm.Data()
+    inp2.store()
+
+    # Create calc with inputs
+    calc = orm.CalculationNode()
+    calc.add_incoming(inp1, link_type=LinkType.INPUT_CALC, link_label='inp1label')
+    calc.add_incoming(inp2, link_type=LinkType.INPUT_CALC, link_label='inp2label')
+    calc.store()
+
+    # Attach outputs
+    out1 = orm.Data()
+    out2 = orm.Data()
+    out1.add_incoming(calc, link_type=LinkType.CREATE, link_label='out1label')
+    out1.store()
+    out2.add_incoming(calc, link_type=LinkType.CREATE, link_label='out2label')
+    out2.store()
+
+    expected_inputs = {'inp1label': inp1.uuid, 'inp2label': inp2.uuid}
+    expected_outputs = {'out1label': out1.uuid, 'out2label': out2.uuid}
+
+    #### Check the 'inputs' manager ###
+    # Check that dir() return all keys and nothing else, important
+    # for tab competion (we skip anything that starts with an underscore)
+    assert len([key for key in dir(calc.inputs) if not key.startswith('_')]) == len(expected_inputs)
+    assert set(key for key in dir(calc.inputs) if not key.startswith('_')) == set(expected_inputs)
+    # Check that it works also as an iterator
+    assert len(list(calc.inputs)) == len(expected_inputs)
+    assert set(calc.inputs) == set(expected_inputs)
+
+    for key, val in expected_inputs.items():
+        # calc.inputs.a.uuid == ..., ...
+        assert getattr(calc.inputs, key).uuid == val
+        # calc.inputs['a'].uuid == ..., ...
+        assert calc.inputs[key].uuid == val
+
+    # I check the attribute fetching directly
+    assert calc.inputs.inp1label.uuid == expected_inputs['inp1label']
+
+    ## Check for not-existing links
+    # - Must raise a AttributeError, otherwise tab competion will not work
+    # - Actually raises a NotExistentAttributeError
+    # - NotExistentAttributeError should also be caught by NotExistent,
+    #   for backwards-compatibility for AiiDA 1.0, 1.1, 1.2
+    for exception in [AttributeError, NotExistent, NotExistentAttributeError]:
+        with pytest.raises(exception):
+            getattr(calc.inputs, 'NotExistentLabel')
+
+    # - Must raise a KeyError to behave like a dictionary
+    # - Actually raises a NotExistentKeyError
+    # - NotExistentKeyError should also be caught by NotExistent,
+    #   for backwards-compatibility for AiiDA 1.0, 1.1, 1.2
+    for exception in [KeyError, NotExistent, NotExistentKeyError]:
+        with pytest.raises(exception):
+            _ = calc.inputs['NotExistentLabel']
+
+    #### Check the 'outputs' manager ###
+    # Check that dir() return all keys and nothing else, important
+    # for tab competion (we skip anything that starts with an underscore)
+    assert len([key for key in dir(calc.outputs) if not key.startswith('_')]) == len(expected_outputs)
+    assert set(key for key in dir(calc.outputs) if not key.startswith('_')) == set(expected_outputs)
+    # Check that it works also as an iterator
+    assert len(list(calc.outputs)) == len(expected_outputs)
+    assert set(calc.outputs) == set(expected_outputs)
+
+    for key, val in expected_outputs.items():
+        # calc.outputs.a.uuid == ..., ...
+        assert getattr(calc.outputs, key).uuid == val
+        # calc.outputs['a'].uuid == ..., ...
+        assert calc.outputs[key].uuid == val
+
+    # I check the attribute fetching directly
+    assert calc.outputs.out1label.uuid == expected_outputs['out1label']
+
+    # Must raise a AttributeError, otherwise tab competion will not work
+    with pytest.raises(AttributeError):
+        getattr(calc.outputs, 'NotExistentLabel')
+
+    # Must raise a KeyError
+    with pytest.raises(KeyError):
+        _ = calc.outputs['NotExistentLabel']