isolated bandsdata list and test

aiida/backends/djsite/queries.py

@@ -108,43 +108,6 @@ def query_group(q_object, args):
         if args.group_pk is not None:
             q_object.add(Q(dbgroups__pk__in=args.group_pk), Q.AND)
 
-    @staticmethod
-    def _extract_formula(struc_pk, args, deser_data):
-        """Extract formula."""
-        from aiida.orm.nodes.data.structure import (get_formula, get_symbols_string)
-
-        if struc_pk is not None:
-            # Exclude structures by the elements
-            if args.element is not None:
-                all_kinds = [k['symbols'] for k in deser_data[struc_pk]['kinds']]
-                all_symbols = [item for sublist in all_kinds for item in sublist]
-                if not any([s in args.element for s in all_symbols]):
-                    return None
-            if args.element_only is not None:
-                all_kinds = [k['symbols'] for k in deser_data[struc_pk]['kinds']]
-                all_symbols = [item for sublist in all_kinds for item in sublist]
-                if not all([s in all_symbols for s in args.element_only]):
-                    return None
-
-            # build the formula
-            symbol_dict = {
-                k['name']: get_symbols_string(k['symbols'], k['weights']) for k in deser_data[struc_pk]['kinds']
-            }
-            try:
-                symbol_list = [symbol_dict[s['kind_name']] for s in deser_data[struc_pk]['sites']]
-                formula = get_formula(symbol_list, mode=args.formulamode)
-            # If for some reason there is no kind with the name
-            # referenced by the site
-            except KeyError:
-                formula = '<<UNKNOWN>>'
-                # cycle if we imposed the filter on elements
-                if args.element is not None or args.element_only is not None:
-                    return None
-        else:
-            formula = '<<UNKNOWN>>'
-
-        return formula
-
     def get_bands_and_parents_structure(self, args):
         """Returns bands and closest parent structure."""
         from django.db.models import Q
@@ -175,14 +138,24 @@ def get_bands_and_parents_structure(self, args):
             # get the closest structures (WITHOUT DbPath)
             structure_dict = get_closest_parents(pks, Q(node_type='data.structure.StructureData.'), chunk_size=1)
 
-            struc_pks = [structure_dict[pk] for pk in pks]
+            struc_pks = [structure_dict.get(pk) for pk in pks]
 
             # query for the attributes needed for the structure formula
             res_attr = models.DbNode.objects.filter(id__in=struc_pks).values_list('id', 'attributes')
+            res_attr = {rattr[0]: rattr[1] for rattr in res_attr}
 
             # prepare the printout
             for (b_id_lbl_date, struc_pk) in zip(this_chunk, struc_pks):
-                formula = self._extract_formula(struc_pk, args, {rattr[0]: rattr[1] for rattr in res_attr})
+                if struc_pk is not None:
+                    strct = res_attr[struc_pk]
+                    akinds, asites = strct['kinds'], strct['sites']
+                    formula = self._extract_formula(akinds, asites, args)
+                else:
+                    if args.element is not None or args.element_only is not None:
+                        formula = None
+                    else:
+                        formula = '<<NOT FOUND>>'
+
                 if formula is None:
                     continue
                 entry_list.append([

aiida/backends/general/abstractqueries.py

@@ -120,8 +120,19 @@ def get_statistics_dict(dataset):
         return statistics
 
     @staticmethod
-    def _extract_formula(args, akinds, asites):
-        """Extract formula from the structure object."""
+    def _extract_formula(akinds, asites, args):
+        """
+        Extract formula from the structure object.
+
+        :param akinds: list of kinds, e.g. [{'mass': 55.845, 'name': 'Fe', 'symbols': ['Fe'], 'weights': [1.0]},
+                                            {'mass': 15.9994, 'name': 'O', 'symbols': ['O'], 'weights': [1.0]}]
+        :param asites: list of structure sites e.g. [{'position': [0.0, 0.0, 0.0], 'kind_name': 'Fe'},
+                                                        {'position': [2.0, 2.0, 2.0], 'kind_name': 'O'}]
+        :param args: a namespace with parsed command line parameters, here only 'element' and 'element_only' are used
+        :type args: dict
+
+        :return: a string with formula if the formula is found
+        """
         from aiida.orm.nodes.data.structure import (get_formula, get_symbols_string)
 
         if args.element is not None:
@@ -136,7 +147,7 @@ def _extract_formula(args, akinds, asites):
 
         # We want only the StructureData that have attributes
         if akinds is None or asites is None:
-            return None
+            return '<<UNKNOWN>>'
 
         symbol_dict = {}
         for k in akinds:
@@ -161,7 +172,9 @@ def get_bands_and_parents_structure(self, args):
 
         :returns:
             A list of sublists, each latter containing (in order):
-                pk as string, formula as string, creation date, bandsdata-label"""
+                pk as string, formula as string, creation date, bandsdata-label
+        """
+        # pylint: disable=too-many-locals
 
         import datetime
         from aiida.common import timezone
@@ -173,22 +186,23 @@ def get_bands_and_parents_structure(self, args):
         else:
             q_build.append(orm.User, tag='creator')
 
-        bdata_filters = {}
-        if args.past_days is not None:
-            bdata_filters.update({'ctime': {'>=': timezone.now() - datetime.timedelta(days=args.past_days)}})
-
-        q_build.append(
-            orm.BandsData, tag='bdata', with_user='creator', filters=bdata_filters, project=['id', 'label', 'ctime']
-        )
-
         group_filters = {}
 
         if args.group_name is not None:
             group_filters.update({'name': {'in': args.group_name}})
         if args.group_pk is not None:
             group_filters.update({'id': {'in': args.group_pk}})
-        if group_filters:
-            q_build.append(orm.Group, tag='group', filters=group_filters, with_node='bdata')
+
+        q_build.append(orm.Group, tag='group', filters=group_filters, with_user='creator')
+
+        bdata_filters = {}
+        if args.past_days is not None:
+            bdata_filters.update({'ctime': {'>=': timezone.now() - datetime.timedelta(days=args.past_days)}})
+
+        q_build.append(
+            orm.BandsData, tag='bdata', with_group='group', filters=bdata_filters, project=['id', 'label', 'ctime']
+        )
+        bands_list_data = q_build.all()
 
         q_build.append(
             orm.StructureData,
@@ -200,12 +214,15 @@ def get_bands_and_parents_structure(self, args):
 
         q_build.order_by({orm.StructureData: {'ctime': 'desc'}})
 
-        list_data = q_build.distinct()
+        structure_dict = dict()
+        list_data = q_build.distinct().all()
+        for bid, _, _, _, akinds, asites in list_data:
+            structure_dict[bid] = (akinds, asites)
 
         entry_list = []
         already_visited_bdata = set()
 
-        for [bid, blabel, bdate, _, akinds, asites] in list_data.all():
+        for [bid, blabel, bdate] in bands_list_data:
 
             # We process only one StructureData per BandsData.
             # We want to process the closest StructureData to
@@ -217,7 +234,17 @@ def get_bands_and_parents_structure(self, args):
             if already_visited_bdata.__contains__(bid):
                 continue
             already_visited_bdata.add(bid)
-            formula = self._extract_formula(args, akinds, asites)
+            strct = structure_dict.get(bid, None)
+
+            if strct is not None:
+                akinds, asites = strct
+                formula = self._extract_formula(akinds, asites, args)
+            else:
+                if args.element is not None or args.element_only is not None:
+                    formula = None
+                else:
+                    formula = '<<NOT FOUND>>'
+
             if formula is None:
                 continue
             entry_list.append([str(bid), str(formula), bdate.strftime('%d %b %Y'), blabel])

tests/cmdline/commands/test_data.py

@@ -141,7 +141,6 @@ def data_listing_test(self, datatype, search_string, ids):
 
         # Check that the past days filter works as expected
         past_days_flags = ['-p', '--past-days']
-        # past_days_flags = ['-p']
         for flag in past_days_flags:
             options = [flag, '1']
             res = self.cli_runner.invoke(listing_cmd, options, catch_exceptions=False)
@@ -158,6 +157,7 @@ def data_listing_test(self, datatype, search_string, ids):
             )
 
         # Check that the group filter works as expected
+        # if ids is not None:
         group_flags = ['-G', '--groups']
         for flag in group_flags:
             # Non empty group
@@ -289,10 +289,14 @@ def connect_structure_bands(strct):  # pylint: disable=unused-argument
 
         bands = connect_structure_bands(strct)
 
+        bands_isolated = BandsData()
+        bands_isolated.store()
+
         # Create 2 groups and add the data to one of them
         g_ne = Group(label='non_empty_group')
         g_ne.store()
         g_ne.add_nodes(bands)
+        g_ne.add_nodes(bands_isolated)
 
         g_e = Group(label='empty_group')
         g_e.store()
@@ -321,6 +325,13 @@ def test_bandlistshelp(self):
 
     def test_bandslist(self):
         self.data_listing_test(BandsData, 'FeO', self.ids)
+        self.data_listing_test(BandsData, '<<NOT FOUND>>', self.ids)
+
+    def test_bandslist_with_elements(self):
+        options = ['-e', 'Fe']
+        res = self.cli_runner.invoke(cmd_bands.bands_list, options, catch_exceptions=False)
+        self.assertIn(b'FeO', res.stdout_bytes, 'The string "FeO" was not found in the listing')
+        self.assertNotIn(b'<<NOT FOUND>>', res.stdout_bytes, 'The string "<<NOT FOUND>>" should not in the listing')
 
     def test_bandexporthelp(self):
         output = sp.check_output(['verdi', 'data', 'bands', 'export', '--help'])