Skip to content

Commit

Permalink
ORM: Add the Entity.get_collection classmethod
Browse files Browse the repository at this point in the history
This is an alternative for the `collection` class property. Where the
`collection` property uses the current default storage backend, the
`get_collection` allows to specify another specific backend.

The original design intended to support this use-case by allowing a
`Collection` instance to be "called" with a specific backend:

    Entity.collection(backend)

The `.collection` returns a `Collection` instance which is then called
with passing the `backend`, which would return a new `Collection` for
the same entity type, but with the other backend. However, this doesn't
always work, because the `collection` property will load the backend
from the default profile, which will except if not loaded. Although this
scenario is unlikely for normal usage, application developers may use
AiiDA's API with multiple active backends where no default profile is
defined.

The reason for a new method instead of changing the `collection`
property is that that would be backwards incompatible.
  • Loading branch information
sphuber committed Sep 5, 2023
1 parent 2071517 commit 305f1db
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
12 changes: 12 additions & 0 deletions aiida/orm/entities.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,18 @@ def collection(cls) -> CollectionType: # pylint: disable=no-self-argument
"""
return cls._CLS_COLLECTION.get_cached(cls, get_manager().get_profile_storage())

@classmethod
def get_collection(cls, backend: 'StorageBackend'):
"""Get a collection for objects of this type for a given backend.
.. note:: Use the ``collection`` class property instead if the currently loaded backend or backend of the
default profile should be used.
:param backend: The backend of the collection to use.
:return: A collection object that can be used to access entities of this type.
"""
return cls._CLS_COLLECTION.get_cached(cls, backend)

@classmethod
def get(cls, **kwargs):
"""Get an entity of the collection matching the given filters.
Expand Down
6 changes: 6 additions & 0 deletions tests/orm/test_entities.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,12 @@
class TestBackendEntitiesAndCollections:
"""Test backend entities and their collections"""

def test_get_collection(self, backend):
"""Test :meth:`aiida.orm.entities.Entity.get_collection`."""
user_collection = orm.User.get_collection(backend)
assert user_collection is orm.User.collection
assert user_collection.backend is backend

def test_collections_cache(self):
"""Make sure that we're not recreating collections each time .collection is called"""
# Check directly
Expand Down

0 comments on commit 305f1db

Please sign in to comment.