diff --git a/aiida/backends/djsite/queries.py b/aiida/backends/djsite/queries.py index 209d646306..ff8121c7ab 100644 --- a/aiida/backends/djsite/queries.py +++ b/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 = '<>' - # 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 = '<>' - - 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 = '<>' + if formula is None: continue entry_list.append([ diff --git a/aiida/backends/general/abstractqueries.py b/aiida/backends/general/abstractqueries.py index bf8a74bca1..1851eca2c6 100644 --- a/aiida/backends/general/abstractqueries.py +++ b/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 '<>' 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 = '<>' + if formula is None: continue entry_list.append([str(bid), str(formula), bdate.strftime('%d %b %Y'), blabel]) diff --git a/tests/cmdline/commands/test_data.py b/tests/cmdline/commands/test_data.py index ff63bfa927..11764640e3 100644 --- a/tests/cmdline/commands/test_data.py +++ b/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, '<>', 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'<>', res.stdout_bytes, 'The string "<>" should not in the listing') def test_bandexporthelp(self): output = sp.check_output(['verdi', 'data', 'bands', 'export', '--help'])