Skip to content

Commit

Permalink
♻️ REFACTOR: EntityAttributesMixin -> NodeAttributes (aiidateam#5442
Browse files Browse the repository at this point in the history 
)

This commit is part of the `Node` namespace restructure (aiidateam#5465).
It moves all attribute related methods to a new `NodeAttributes` class,
accessed via `Node.base.attributes`.

The commit also fixes a bug with the `Sealable` mixin of `ProcessNode`,
whereby it only overrode `set_attribute` and `delete_attribute`,
to change mutability behaviour, but not any other mutation methods, like `set_attribute_many`, etc
  • Loading branch information
@chrisjsewell
chrisjsewell authored Apr 8, 2022
1 parent 8293e45 commit 8bb6c3f
Show file tree
Hide file tree
Showing 67 changed files with 767 additions and 720 deletions.
2 changes: 1 addition & 1 deletion .molecule/default/files/polish/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ def run_via_daemon(workchains, inputs, sleep, timeout):
except AttributeError:
click.secho('Failed: ', fg='red', bold=True, nl=False)
click.secho(f'the workchain<{workchain.pk}> did not return a result output node', bold=True)
click.echo(str(workchain.attributes))
click.echo(str(workchain.base.attributes.all))
return None

return result, workchain, total_time
Expand Down
1 change: 1 addition & 0 deletions .pre-commit-config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ repos:
aiida/orm/users.py|
aiida/orm/nodes/data/enum.py|
aiida/orm/nodes/data/jsonable.py|
aiida/orm/nodes/attributes.py|
aiida/orm/nodes/node.py|
aiida/orm/nodes/process/.*py|
aiida/orm/nodes/repository.py|
Expand Down
2 changes: 1 addition & 1 deletion aiida/cmdline/commands/cmd_node.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -234,7 +234,7 @@ def echo_node_dict(nodes, keys, fmt, identifier, raw, use_attrs=True):
id_value = node.uuid

if use_attrs:
node_dict = node.attributes
node_dict = node.base.attributes.all
dict_name = 'attributes'
else:
node_dict = node.extras
Expand Down
5 changes: 4 additions & 1 deletion aiida/cmdline/utils/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,10 @@ def format_flat_links(links, headers):
table = []

for link_triple in links:
table.append([link_triple.link_label, link_triple.node.pk, link_triple.node.get_attribute('process_label', '')])
table.append([
link_triple.link_label, link_triple.node.pk,
link_triple.node.base.attributes.get('process_label', '')
])

result = f'\n{tabulate(table, headers=headers)}'

Expand Down
2 changes: 1 addition & 1 deletion aiida/engine/processes/calcjobs/calcjob.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -440,7 +440,7 @@ def _perform_import(self):
)
retrieve_calculation(self.node, transport, retrieved_temporary_folder.abspath)
self.node.set_state(CalcJobState.PARSING)
self.node.set_attribute(orm.CalcJobNode.IMMIGRATED_KEY, True)
self.node.base.attributes.set(orm.CalcJobNode.IMMIGRATED_KEY, True)
return self.parse(retrieved_temporary_folder.abspath)

def parse(self, retrieved_temporary_folder: Optional[str] = None) -> ExitCode:
Expand Down
2 changes: 1 addition & 1 deletion aiida/engine/processes/process.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -685,7 +685,7 @@ def _setup_db_record(self) -> None:
def _setup_metadata(self) -> None:
"""Store the metadata on the ProcessNode."""
version_info = self.runner.plugin_version_provider.get_version_info(self.__class__)
self.node.set_attribute_many(version_info)
self.node.base.attributes.set_many(version_info)

for name, metadata in self.metadata.items():
if name in ['store_provenance', 'dry_run', 'call_link_label']:
Expand Down
2 changes: 1 addition & 1 deletion aiida/orm/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,6 @@
'Data',
'Dict',
'Entity',
'EntityAttributesMixin',
'EntityExtrasMixin',
'EntityTypes',
'EnumData',
Expand All @@ -70,6 +69,7 @@
'List',
'Log',
'Node',
'NodeAttributes',
'NodeEntityLoader',
'NodeLinksManager',
'NodeRepository',
Expand Down
162 changes: 1 addition & 161 deletions aiida/orm/entities.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,14 @@

from plumpy.base.utils import call_with_super_check, super_check

from aiida.common import exceptions
from aiida.common.lang import classproperty, type_check
from aiida.manage import get_manager

if TYPE_CHECKING:
from aiida.orm.implementation import BackendEntity, StorageBackend
from aiida.orm.querybuilder import FilterType, OrderByType, QueryBuilder

__all__ = ('Entity', 'Collection', 'EntityAttributesMixin', 'EntityExtrasMixin', 'EntityTypes')
__all__ = ('Entity', 'Collection', 'EntityExtrasMixin', 'EntityTypes')

CollectionType = TypeVar('CollectionType', bound='Collection')
EntityType = TypeVar('EntityType', bound='Entity')
Expand Down Expand Up @@ -260,165 +259,6 @@ def is_stored(self) -> bool:
...


class EntityAttributesMixin:
"""Mixin class that adds all methods for the attributes column to an entity."""

@property
def attributes(self: EntityProtocol) -> Dict[str, Any]:
"""Return the complete attributes dictionary.
.. warning:: While the entity is unstored, this will return references of the attributes on the database model,
meaning that changes on the returned values (if they are mutable themselves, e.g. a list or dictionary) will
automatically be reflected on the database model as well. As soon as the entity is stored, the returned
attributes will be a deep copy and mutations of the database attributes will have to go through the
appropriate set methods. Therefore, once stored, retrieving a deep copy can be a heavy operation. If you
only need the keys or some values, use the iterators `attributes_keys` and `attributes_items`, or the
getters `get_attribute` and `get_attribute_many` instead.
:return: the attributes as a dictionary
"""
attributes = self.backend_entity.attributes

if self.is_stored:
attributes = copy.deepcopy(attributes)

return attributes

def get_attribute(self: EntityProtocol, key: str, default=_NO_DEFAULT) -> Any:
"""Return the value of an attribute.
.. warning:: While the entity is unstored, this will return a reference of the attribute on the database model,
meaning that changes on the returned value (if they are mutable themselves, e.g. a list or dictionary) will
automatically be reflected on the database model as well. As soon as the entity is stored, the returned
attribute will be a deep copy and mutations of the database attributes will have to go through the
appropriate set methods.
:param key: name of the attribute
:param default: return this value instead of raising if the attribute does not exist
:return: the value of the attribute
:raises AttributeError: if the attribute does not exist and no default is specified
"""
try:
attribute = self.backend_entity.get_attribute(key)
except AttributeError:
if default is _NO_DEFAULT:
raise
attribute = default

if self.is_stored:
attribute = copy.deepcopy(attribute)

return attribute

def get_attribute_many(self: EntityProtocol, keys: List[str]) -> List[Any]:
"""Return the values of multiple attributes.
.. warning:: While the entity is unstored, this will return references of the attributes on the database model,
meaning that changes on the returned values (if they are mutable themselves, e.g. a list or dictionary) will
automatically be reflected on the database model as well. As soon as the entity is stored, the returned
attributes will be a deep copy and mutations of the database attributes will have to go through the
appropriate set methods. Therefore, once stored, retrieving a deep copy can be a heavy operation. If you
only need the keys or some values, use the iterators `attributes_keys` and `attributes_items`, or the
getters `get_attribute` and `get_attribute_many` instead.
:param keys: a list of attribute names
:return: a list of attribute values
:raises AttributeError: if at least one attribute does not exist
"""
attributes = self.backend_entity.get_attribute_many(keys)

if self.is_stored:
attributes = copy.deepcopy(attributes)

return attributes

def set_attribute(self: EntityProtocol, key: str, value: Any) -> None:
"""Set an attribute to the given value.
:param key: name of the attribute
:param value: value of the attribute
:raise aiida.common.ValidationError: if the key is invalid, i.e. contains periods
:raise aiida.common.ModificationNotAllowed: if the entity is stored
"""
if self.is_stored:
raise exceptions.ModificationNotAllowed('the attributes of a stored entity are immutable')

self.backend_entity.set_attribute(key, value)

def set_attribute_many(self: EntityProtocol, attributes: Dict[str, Any]) -> None:
"""Set multiple attributes.
.. note:: This will override any existing attributes that are present in the new dictionary.
:param attributes: a dictionary with the attributes to set
:raise aiida.common.ValidationError: if any of the keys are invalid, i.e. contain periods
:raise aiida.common.ModificationNotAllowed: if the entity is stored
"""
if self.is_stored:
raise exceptions.ModificationNotAllowed('the attributes of a stored entity are immutable')

self.backend_entity.set_attribute_many(attributes)

def reset_attributes(self: EntityProtocol, attributes: Dict[str, Any]) -> None:
"""Reset the attributes.
.. note:: This will completely clear any existing attributes and replace them with the new dictionary.
:param attributes: a dictionary with the attributes to set
:raise aiida.common.ValidationError: if any of the keys are invalid, i.e. contain periods
:raise aiida.common.ModificationNotAllowed: if the entity is stored
"""
if self.is_stored:
raise exceptions.ModificationNotAllowed('the attributes of a stored entity are immutable')

self.backend_entity.reset_attributes(attributes)

def delete_attribute(self: EntityProtocol, key: str) -> None:
"""Delete an attribute.
:param key: name of the attribute
:raises AttributeError: if the attribute does not exist
:raise aiida.common.ModificationNotAllowed: if the entity is stored
"""
if self.is_stored:
raise exceptions.ModificationNotAllowed('the attributes of a stored entity are immutable')

self.backend_entity.delete_attribute(key)

def delete_attribute_many(self: EntityProtocol, keys: List[str]) -> None:
"""Delete multiple attributes.
:param keys: names of the attributes to delete
:raises AttributeError: if at least one of the attribute does not exist
:raise aiida.common.ModificationNotAllowed: if the entity is stored
"""
if self.is_stored:
raise exceptions.ModificationNotAllowed('the attributes of a stored entity are immutable')

self.backend_entity.delete_attribute_many(keys)

def clear_attributes(self: EntityProtocol) -> None:
"""Delete all attributes."""
if self.is_stored:
raise exceptions.ModificationNotAllowed('the attributes of a stored entity are immutable')

self.backend_entity.clear_attributes()

def attributes_items(self: EntityProtocol):
"""Return an iterator over the attributes.
:return: an iterator with attribute key value pairs
"""
return self.backend_entity.attributes_items()

def attributes_keys(self: EntityProtocol):
"""Return an iterator over the attribute keys.
:return: an iterator with attribute keys
"""
return self.backend_entity.attributes_keys()


class EntityExtrasMixin:
"""Mixin class that adds all methods for the extras column to an entity."""

Expand Down
2 changes: 2 additions & 0 deletions aiida/orm/nodes/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# yapf: disable
# pylint: disable=wildcard-import

from .attributes import *
from .data import *
from .node import *
from .process import *
Expand All @@ -40,6 +41,7 @@
'KpointsData',
'List',
'Node',
'NodeAttributes',
'NodeRepository',
'NumericType',
'OrbitalData',
Expand Down
Loading

0 comments on commit 8bb6c3f

Please sign in to comment.