From 305f1dbf4ccb3e0e2e79865aee8d248e5ad55b95 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Fri, 1 Sep 2023 15:23:47 +0200 Subject: [PATCH] ORM: Add the `Entity.get_collection` classmethod 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. --- aiida/orm/entities.py | 12 ++++++++++++ tests/orm/test_entities.py | 6 ++++++ 2 files changed, 18 insertions(+) diff --git a/aiida/orm/entities.py b/aiida/orm/entities.py index cec8122d88..fdf0639db2 100644 --- a/aiida/orm/entities.py +++ b/aiida/orm/entities.py @@ -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. diff --git a/tests/orm/test_entities.py b/tests/orm/test_entities.py index ae9b1272d4..bbdc4fa5cf 100644 --- a/tests/orm/test_entities.py +++ b/tests/orm/test_entities.py @@ -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