# -*- 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 #
###########################################################################
"""Module for all common top level AiiDA entity classes and methods"""
import abc
from enum import Enum
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Generic, List, Optional, Type, TypeVar, cast
from plumpy.base.utils import call_with_super_check, super_check
from aiida.common.lang import classproperty, type_check
from aiida.common.warnings import warn_deprecation
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', 'EntityTypes')
CollectionType = TypeVar('CollectionType', bound='Collection')
EntityType = TypeVar('EntityType', bound='Entity')
BackendEntityType = TypeVar('BackendEntityType', bound='BackendEntity')
[docs]class EntityTypes(Enum):
"""Enum for referring to ORM entities in a backend-agnostic manner."""
AUTHINFO = 'authinfo'
COMMENT = 'comment'
COMPUTER = 'computer'
GROUP = 'group'
LOG = 'log'
NODE = 'node'
USER = 'user'
LINK = 'link'
GROUP_NODE = 'group_node'
[docs]class Collection(abc.ABC, Generic[EntityType]):
"""Container class that represents the collection of objects of a particular entity type."""
[docs] @staticmethod
@abc.abstractmethod
def _entity_base_cls() -> Type[EntityType]:
"""The allowed entity class or subclasses thereof."""
[docs] @classmethod
@lru_cache(maxsize=100)
def get_cached(cls, entity_class: Type[EntityType], backend: 'StorageBackend'):
"""Get the cached collection instance for the given entity class and backend.
:param backend: the backend instance to get the collection for
"""
from aiida.orm.implementation import StorageBackend
type_check(backend, StorageBackend)
return cls(entity_class, backend=backend)
[docs] def __init__(self, entity_class: Type[EntityType], backend: Optional['StorageBackend'] = None) -> None:
""" Construct a new entity collection.
:param entity_class: the entity type e.g. User, Computer, etc
:param backend: the backend instance to get the collection for, or use the default
"""
from aiida.orm.implementation import StorageBackend
type_check(backend, StorageBackend, allow_none=True)
assert issubclass(entity_class, self._entity_base_cls())
self._backend = backend or get_manager().get_profile_storage()
self._entity_type = entity_class
[docs] def __call__(self: CollectionType, backend: 'StorageBackend') -> CollectionType:
"""Get or create a cached collection using a new backend."""
if backend is self._backend:
return self
return self.get_cached(self.entity_type, backend=backend) # type: ignore
@property
def entity_type(self) -> Type[EntityType]:
"""The entity type for this instance."""
return self._entity_type
@property
def backend(self) -> 'StorageBackend':
"""Return the backend."""
return self._backend
[docs] def query(
self,
filters: Optional['FilterType'] = None,
order_by: Optional['OrderByType'] = None,
limit: Optional[int] = None,
offset: Optional[int] = None
) -> 'QueryBuilder':
"""Get a query builder for the objects of this collection.
:param filters: the keyword value pair filters to match
:param order_by: a list of (key, direction) pairs specifying the sort order
:param limit: the maximum number of results to return
:param offset: number of initial results to be skipped
"""
from . import querybuilder
filters = filters or {}
order_by = {self.entity_type: order_by} if order_by else {}
query = querybuilder.QueryBuilder(backend=self._backend, limit=limit, offset=offset)
query.append(self.entity_type, project='*', filters=filters)
query.order_by([order_by])
return query
[docs] def get(self, **filters: Any) -> EntityType:
"""Get a single collection entry that matches the filter criteria.
:param filters: the filters identifying the object to get
:return: the entry
"""
res = self.query(filters=filters)
return res.one()[0]
[docs] def find(
self,
filters: Optional['FilterType'] = None,
order_by: Optional['OrderByType'] = None,
limit: Optional[int] = None
) -> List[EntityType]:
"""Find collection entries matching the filter criteria.
:param filters: the keyword value pair filters to match
:param order_by: a list of (key, direction) pairs specifying the sort order
:param limit: the maximum number of results to return
:return: a list of resulting matches
"""
query = self.query(filters=filters, order_by=order_by, limit=limit)
return cast(List[EntityType], query.all(flat=True))
[docs] def all(self) -> List[EntityType]:
"""Get all entities in this collection.
:return: A list of all entities
"""
return cast(List[EntityType], self.query().all(flat=True)) # pylint: disable=no-member
[docs] def count(self, filters: Optional['FilterType'] = None) -> int:
"""Count entities in this collection according to criteria.
:param filters: the keyword value pair filters to match
:return: The number of entities found using the supplied criteria
"""
return self.query(filters=filters).count()
[docs]class Entity(abc.ABC, Generic[BackendEntityType]):
"""An AiiDA entity"""
_CLS_COLLECTION = Collection
@classproperty
def objects(cls: EntityType) -> Collection[EntityType]: # pylint: disable=no-self-argument,no-self-use
"""Get a collection for objects of this type, with the default backend.
.. deprecated:: This will be removed in v3, use ``collection`` instead.
:return: an object that can be used to access entities of this type
"""
warn_deprecation('This property is deprecated, use `.collection` instead.', version=3, stacklevel=2)
return cls.collection
@classproperty
def collection(cls: EntityType) -> Collection[EntityType]: # pylint: disable=no-self-argument,no-self-use
"""Get a collection for objects of this type, with the default backend.
:return: an object that can be used to access entities of this type
"""
return cls._CLS_COLLECTION.get_cached(cls, get_manager().get_profile_storage())
[docs] @classmethod
def get(cls, **kwargs):
"""Get an entity of the collection matching the given filters.
.. deprecated: Will be removed in v3, use `Entity.collection.get` instead.
"""
warn_deprecation(
f'This method is deprecated, use `{cls.__name__}.collection.get` instead.', version=3, stacklevel=2
)
return cls.collection.get(**kwargs) # pylint: disable=no-member
[docs] @classmethod
def from_backend_entity(cls: Type[EntityType], backend_entity: BackendEntityType) -> EntityType:
"""Construct an entity from a backend entity instance
:param backend_entity: the backend entity
:return: an AiiDA entity instance
"""
from .implementation.entities import BackendEntity
type_check(backend_entity, BackendEntity)
entity = cls.__new__(cls)
entity._backend_entity = backend_entity
call_with_super_check(entity.initialize)
return entity
[docs] def __init__(self, backend_entity: BackendEntityType) -> None:
"""
:param backend_entity: the backend model supporting this entity
"""
self._backend_entity = backend_entity
call_with_super_check(self.initialize)
@super_check
def initialize(self) -> None:
"""Initialize instance attributes.
This will be called after the constructor is called or an entity is created from an existing backend entity.
"""
@property
def id(self) -> int: # pylint: disable=invalid-name
"""Return the id for this entity.
This identifier is guaranteed to be unique amongst entities of the same type for a single backend instance.
.. deprecated: Will be removed in v3, use `pk` instead.
:return: the entity's id
"""
warn_deprecation('This method is deprecated, use `pk` instead.', version=3, stacklevel=2)
return self._backend_entity.id
@property
def pk(self) -> int:
"""Return the primary key for this entity.
This identifier is guaranteed to be unique amongst entities of the same type for a single backend instance.
:return: the entity's principal key
"""
return self._backend_entity.id
[docs] def store(self: EntityType) -> EntityType:
"""Store the entity."""
self._backend_entity.store()
return self
@property
def is_stored(self) -> bool:
"""Return whether the entity is stored."""
return self._backend_entity.is_stored
@property
def backend(self) -> 'StorageBackend':
"""Get the backend for this entity"""
return self._backend_entity.backend
@property
def backend_entity(self) -> BackendEntityType:
"""Get the implementing class for this object"""
return self._backend_entity