Adapted the psf and psml data classes in order to support the new Gro…

…up system. From aiida 1.2.0 the class `Group` have been modified to implement new features. This has implications in our plugin because the PseudoFamilies are Groups. In particular there is now an explicit implementation of PsmlFamily and PsfFamily as subclasses of the `Group` class. Moreover PsmlFamily and PsfFamily are registred as entry points making possible to reload the correct sub class when reloading from the database. For more info see: aiidateam/aiida-core#3882
siesta-project · May 19, 2020 · 9de93ed · 9de93ed
1 parent 3beb614
commit 9de93ed
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 56 deletions.
diff --git a/aiida_siesta/commands/data_psf.py b/aiida_siesta/commands/data_psf.py
@@ -55,22 +55,14 @@ def psf_listfamilies(elements, with_description):
     """
     from aiida import orm
     from aiida.plugins import DataFactory
-    from aiida_siesta.data.psf import PSFGROUP_TYPE
+    from aiida_siesta.groups.pseudos import PsfFamily
 
     PsfData = DataFactory('siesta.psf')  # pylint: disable=invalid-name
     query = orm.QueryBuilder()
     query.append(PsfData, tag='psfdata')
     if elements is not None:
         query.add_filter(PsfData, {'attributes.element': {'in': elements}})
-    query.append(
-        orm.Group,
-        with_node='psfdata',
-        tag='group',
-        project=['label', 'description'],
-        filters={'type_string': {
-            '==': PSFGROUP_TYPE
-        }}
-    )
+    query.append(PsfFamily, with_node='psfdata', tag='group', project=['label', 'description'])
 
     query.distinct()
     if query.count() > 0:

diff --git a/aiida_siesta/commands/data_psml.py b/aiida_siesta/commands/data_psml.py
@@ -55,22 +55,14 @@ def psml_listfamilies(elements, with_description):
     """
     from aiida import orm
     from aiida.plugins import DataFactory
-    from aiida_siesta.data.psml import PSMLGROUP_TYPE
+    from aiida_siesta.groups.pseudos import PsmlFamily
 
     PsmlData = DataFactory('siesta.psml')  # pylint: disable=invalid-name
     query = orm.QueryBuilder()
     query.append(PsmlData, tag='psmldata')
     if elements is not None:
         query.add_filter(PsmlData, {'attributes.element': {'in': elements}})
-    query.append(
-        orm.Group,
-        with_node='psmldata',
-        tag='group',
-        project=['label', 'description'],
-        filters={'type_string': {
-            '==': PSMLGROUP_TYPE
-        }}
-    )
+    query.append(PsmlFamily, with_node='psmldata', tag='group', project=['label', 'description'])
 
     query.distinct()
     if query.count() > 0:

diff --git a/aiida_siesta/data/psf.py b/aiida_siesta/data/psf.py
@@ -3,12 +3,9 @@
 """
 
 import io
-from aiida.common.utils import classproperty
 from aiida.common.files import md5_file
 from aiida.orm.nodes import SinglefileData
 
-PSFGROUP_TYPE = 'data.psf.family'
-
 
 def get_pseudos_from_structure(structure, family_name):
     """
@@ -64,6 +61,7 @@ def upload_psf_family(folder, group_label, group_description, stop_if_existing=T
     from aiida.common import AIIDA_LOGGER as aiidalogger
     from aiida.common.exceptions import UniquenessError
     from aiida.orm.querybuilder import QueryBuilder
+    from aiida_siesta.groups.pseudos import PsfFamily
 
     if not os.path.isdir(folder):
         raise ValueError("folder must be a directory")
@@ -79,9 +77,7 @@ def upload_psf_family(folder, group_label, group_description, stop_if_existing=T
     nfiles = len(files)
 
     automatic_user = orm.User.objects.get_default()
-    group, group_created = orm.Group.objects.get_or_create(
-        label=group_label, type_string=PSFGROUP_TYPE, user=automatic_user
-    )
+    group, group_created = PsfFamily.objects.get_or_create(label=group_label, user=automatic_user)
 
     if group.user.email != automatic_user.email:
         raise UniquenessError(
@@ -257,10 +253,6 @@ def get_or_create(cls, filename, use_first=False, store_psf=True):
 
         return (pseudos[0], False)
 
-    @classproperty
-    def psffamily_type_string(cls):  # pylint: disable=no-self-argument,no-self-use
-        return PSFGROUP_TYPE
-
     def store(self, *args, **kwargs):  # pylint: disable=arguments-differ
         """
         Store the node, reparsing the file so that the md5 and the element
@@ -327,11 +319,11 @@ def get_psf_family_names(self):
         """
         Get the list of all psf family names to which the pseudo belongs
         """
-        from aiida.orm import Group
         from aiida.orm import QueryBuilder
+        from aiida_siesta.groups.pseudos import PsfFamily
 
         query = QueryBuilder()
-        query.append(Group, filters={'type_string': {'==': self.psffamily_type_string}}, tag='group', project='label')
+        query.append(PsfFamily, tag='group', project='label')
         query.append(PsfData, filters={'id': {'==': self.id}}, with_group='group')
 
         return [_[0] for _ in query.all()]
@@ -388,9 +380,9 @@ def get_psf_group(cls, group_label):
         """
         Return the PsfFamily group with the given name.
         """
-        from aiida.orm import Group
+        from aiida_siesta.groups.pseudos import PsfFamily
 
-        return Group.get(label=group_label, type_string=cls.psffamily_type_string)
+        return PsfFamily.get(label=group_label)
 
     @classmethod
     def get_psf_groups(cls, filter_elements=None, user=None):
@@ -405,14 +397,12 @@ def get_psf_groups(cls, filter_elements=None, user=None):
                If defined, it should be either a DbUser instance, or a string
                for the username (that is, the user email).
         """
-        from aiida.orm import Group
         from aiida.orm import QueryBuilder
         from aiida.orm import User
+        from aiida_siesta.groups.pseudos import PsfFamily
 
         query = QueryBuilder()
-        filters = {'type_string': {'==': cls.psffamily_type_string}}
-
-        query.append(Group, filters=filters, tag='group', project='*')
+        query.append(PsfFamily, tag='group', project='*')
 
         if user:
             query.append(User, filters={'email': {'==': user}}, with_group='group')
@@ -423,6 +413,6 @@ def get_psf_groups(cls, filter_elements=None, user=None):
         if filter_elements is not None:
             query.append(PsfData, filters={'attributes.element': {'in': filter_elements}}, with_group='group')
 
-        query.order_by({Group: {'id': 'asc'}})
+        query.order_by({PsfFamily: {'id': 'asc'}})
 
         return [_[0] for _ in query.all()]
diff --git a/aiida_siesta/data/psml.py b/aiida_siesta/data/psml.py
@@ -2,13 +2,12 @@
 This module manages the PSML pseudopotentials in the local repository.
 """
 
-from aiida.common.utils import classproperty
 from aiida.common.files import md5_file
 from aiida.orm.nodes import SinglefileData
 
 # See LICENSE.txt and AUTHORS.txt
 
-PSMLGROUP_TYPE = 'data.psml.family'
+#PSMLGROUP_TYPE = 'data.psml.family'
 
 
 def get_pseudos_from_structure(structure, family_name):
@@ -65,6 +64,7 @@ def upload_psml_family(folder, group_label, group_description, stop_if_existing=
     from aiida.common import AIIDA_LOGGER as aiidalogger
     from aiida.common.exceptions import UniquenessError
     from aiida.orm.querybuilder import QueryBuilder
+    from aiida_siesta.groups.pseudos import PsmlFamily
 
     if not os.path.isdir(folder):
         raise ValueError("folder must be a directory")
@@ -80,9 +80,10 @@ def upload_psml_family(folder, group_label, group_description, stop_if_existing=
     nfiles = len(files)
 
     automatic_user = orm.User.objects.get_default()
-    group, group_created = orm.Group.objects.get_or_create(
-        label=group_label, type_string=PSMLGROUP_TYPE, user=automatic_user
-    )
+    #group, group_created = orm.Group.objects.get_or_create(
+    #    label=group_label, type_string=PSMLGROUP_TYPE, user=automatic_user
+    #)
+    group, group_created = PsmlFamily.objects.get_or_create(label=group_label, user=automatic_user)
 
     if group.user.email != automatic_user.email:
         raise UniquenessError(
@@ -258,9 +259,9 @@ def get_or_create(cls, filename, use_first=False, store_psml=True):
 
         return (pseudos[0], False)
 
-    @classproperty
-    def psmlfamily_type_string(cls):  # pylint: disable=no-self-argument,no-self-use
-        return PSMLGROUP_TYPE
+    #@classproperty
+    #def psmlfamily_type_string(cls):  # pylint: disable=no-self-argument,no-self-use
+    #    return PSMLGROUP_TYPE
 
     def store(self, *args, **kwargs):  # pylint: disable=arguments-differ
         """
@@ -328,11 +329,12 @@ def get_psml_family_names(self):
         """
         Get the list of all psml family names to which the pseudo belongs
         """
-        from aiida.orm import Group
+        #from aiida.orm import Group
         from aiida.orm import QueryBuilder
+        from aiida_siesta.groups.pseudos import PsmlFamily
 
         query = QueryBuilder()
-        query.append(Group, filters={'type_string': {'==': self.psmlfamily_type_string}}, tag='group', project='label')
+        query.append(PsmlFamily, tag='group', project='label')
         query.append(PsmlData, filters={'id': {'==': self.id}}, with_group='group')
 
         return [_[0] for _ in query.all()]
@@ -390,9 +392,10 @@ def get_psml_group(cls, group_label):
         """
         Return the PsmlFamily group with the given name.
         """
-        from aiida.orm import Group
+        #from aiida.orm import Group
+        from aiida_siesta.groups.pseudos import PsmlFamily
 
-        return Group.get(label=group_label, type_string=cls.psmlfamily_type_string)
+        return PsmlFamily.get(label=group_label)
 
     @classmethod
     def get_psml_groups(cls, filter_elements=None, user=None):
@@ -407,14 +410,15 @@ def get_psml_groups(cls, filter_elements=None, user=None):
                If defined, it should be either a DbUser instance, or a string
                for the username (that is, the user email).
         """
-        from aiida.orm import Group
+        #from aiida.orm import Group
         from aiida.orm import QueryBuilder
         from aiida.orm import User
+        from aiida_siesta.groups.pseudos import PsmlFamily
 
         query = QueryBuilder()
-        filters = {'type_string': {'==': cls.psmlfamily_type_string}}
+        #filters = {'type_string': {'==': cls.psmlfamily_type_string}}
 
-        query.append(Group, filters=filters, tag='group', project='*')
+        query.append(PsmlFamily, tag='group', project='*')
 
         if user:
             query.append(User, filters={'email': {'==': user}}, with_group='group')
@@ -425,6 +429,6 @@ def get_psml_groups(cls, filter_elements=None, user=None):
         if filter_elements is not None:
             query.append(PsmlData, filters={'attributes.element': {'in': filter_elements}}, with_group='group')
 
-        query.order_by({Group: {'id': 'asc'}})
+        query.order_by({PsmlFamily: {'id': 'asc'}})
 
         return [_[0] for _ in query.all()]
diff --git a/aiida_siesta/groups/pseudos.py b/aiida_siesta/groups/pseudos.py
@@ -0,0 +1,8 @@
+from aiida.orm.groups import Group
+
+class PsfFamily(Group):
+    """Group that represents a pseudo potential family containing `PsfData` nodes."""
+
+
+class PsmlFamily(Group):
+    """Group that represents a pseudo potential family containing `PsmlData` nodes."""
diff --git a/setup.json b/setup.json
@@ -55,6 +55,10 @@
         "aiida.cmdline.data": [
             "psf = aiida_siesta.commands.data_psf:psfdata",
             "psml = aiida_siesta.commands.data_psml:psmldata"
-        ]
+        ],
+	"aiida.groups": [
+      	    "data.psf.family = aiida_siesta.groups.pseudos:PsfFamily",
+            "data.psml.family = aiida_siesta.groups.pseudos:PsmlFamily"
+    	]
     }
 }