From ecc9bf1c86fd01e9b9fff80fee1e136ffc21ac3d Mon Sep 17 00:00:00 2001 From: Casper Welzel Andersen Date: Sat, 9 Jan 2021 03:23:22 +0100 Subject: [PATCH] Ensure minimized fields works Build attributes only with supplied "desired" fields. This ensures the `-m/--minimized-fields` option for `aiida-optimade init` works as intended. --- aiida_optimade/cli/cmd_init.py | 17 ++++++----- aiida_optimade/entry_collections.py | 9 ++++-- aiida_optimade/mappers/entries.py | 12 +++++++- aiida_optimade/mappers/structures.py | 42 +++++++++++++++++++++------- aiida_optimade/models/structures.py | 11 +++++--- 5 files changed, 67 insertions(+), 24 deletions(-) diff --git a/aiida_optimade/cli/cmd_init.py b/aiida_optimade/cli/cmd_init.py index d8dd1495..70059061 100644 --- a/aiida_optimade/cli/cmd_init.py +++ b/aiida_optimade/cli/cmd_init.py @@ -91,14 +91,17 @@ def init(obj: dict, force: bool, silent: bool, minimized_fields: bool): STRUCTURES._filter_fields = set() if minimized_fields: - STRUCTURES._alias_filter( - dict.fromkeys( - [ - "structure_features", # required (will create species) - ], - None, - ) + minimized_keys = ( + STRUCTURES.resource_mapper.TOP_LEVEL_NON_ATTRIBUTES_FIELDS.copy() + ) + minimized_keys |= STRUCTURES.get_attribute_fields() + minimized_keys |= { + f"_{STRUCTURES.provider}_" + _ for _ in STRUCTURES.provider_fields + } + minimized_keys.difference_update( + {"cartesian_site_positions", "nsites", "species_at_sites"} ) + STRUCTURES._alias_filter(dict.fromkeys(minimized_keys, None)) else: STRUCTURES._alias_filter({"nsites": None}) diff --git a/aiida_optimade/entry_collections.py b/aiida_optimade/entry_collections.py index ab7b3498..f7c283b9 100644 --- a/aiida_optimade/entry_collections.py +++ b/aiida_optimade/entry_collections.py @@ -153,7 +153,7 @@ def find( all_fields = criteria.pop("fields") if getattr(params, "response_fields", False): fields = set(params.response_fields.split(",")) - fields |= self.resource_mapper.get_required_fields() + fields |= self.resource_mapper.TOP_LEVEL_NON_ATTRIBUTES_FIELDS else: fields = all_fields.copy() @@ -362,7 +362,7 @@ def _update_entities(entities: list, fields: list): necessary_entity_ids = [pk[0] for pk in necessary_entities_qb] # Create the missing OPTIMADE fields: - fields = {"id", "type"} + fields = self.resource_mapper.TOP_LEVEL_NON_ATTRIBUTES_FIELDS.copy() if all_fields: # All OPTIMADE fields fields |= self.get_attribute_fields() @@ -373,6 +373,11 @@ def _update_entities(entities: list, fields: list): # "id" and "type" are ALWAYS needed though, hence `fields` is initiated # with these values fields |= self._get_extras_filter_fields() + fields |= { + f"_{self.provider}_" + _ + for _ in self._filter_fields + if _ in self.provider_fields + } fields = list({self.resource_mapper.alias_for(f) for f in fields}) entities = self._find_all( diff --git a/aiida_optimade/mappers/entries.py b/aiida_optimade/mappers/entries.py index e29b231c..63cb8ce0 100644 --- a/aiida_optimade/mappers/entries.py +++ b/aiida_optimade/mappers/entries.py @@ -16,7 +16,6 @@ class ResourceMapper(OptimadeResourceMapper): TRANSLATORS: Dict[str, AiidaEntityTranslator] ALL_ATTRIBUTES: set = set() - REQUIRED_ATTRIBUTES: set = set() @classmethod def all_aliases(cls) -> Tuple[Tuple[str, str]]: @@ -40,6 +39,8 @@ def map_back(cls, entity_properties: dict) -> dict: :return: A resource object in OPTIMADE format :rtype: dict """ + from optimade.server.config import CONFIG + new_object_attributes = {} new_object = {} @@ -64,8 +65,13 @@ def map_back(cls, entity_properties: dict) -> dict: if value is not None: new_object[field] = value + mapping = {aiida: optimade for optimade, aiida in cls.all_aliases()} + new_object["attributes"] = cls.build_attributes( retrieved_attributes=new_object_attributes, + desired_attributes={mapping.get(_, _) for _ in entity_properties} + - cls.TOP_LEVEL_NON_ATTRIBUTES_FIELDS + - set(CONFIG.aliases.get(cls.ENDPOINT, {}).keys()), entry_pk=new_object["id"], node_type=new_object["type"], ) @@ -77,6 +83,7 @@ def map_back(cls, entity_properties: dict) -> dict: def build_attributes( cls, retrieved_attributes: dict, + desired_attributes: list, entry_pk: int, node_type: str, ) -> dict: @@ -85,6 +92,9 @@ def build_attributes( :param retrieved_attributes: Dict of new attributes, will be updated accordingly :type retrieved_attributes: dict + :param desired_attributes: Set of attributes to be built. + :type desired_attributes: set + :param entry_pk: The AiiDA Node's PK :type entry_pk: int diff --git a/aiida_optimade/mappers/structures.py b/aiida_optimade/mappers/structures.py index 87f41e91..ca58ef76 100644 --- a/aiida_optimade/mappers/structures.py +++ b/aiida_optimade/mappers/structures.py @@ -26,21 +26,30 @@ class StructureMapper(ResourceMapper): "data.structure.StructureData.": StructureDataTranslator, } ALL_ATTRIBUTES = set(StructureResourceAttributes.schema().get("properties").keys()) - REQUIRED_ATTRIBUTES = set(StructureResourceAttributes.schema().get("required")) - # This should be REQUIRED_FIELDS, but should be set as such in `optimade` + REQUIRED_FIELDS = set(StructureResourceAttributes.schema().get("required")) # pylint: disable=too-many-locals @classmethod def build_attributes( - cls, retrieved_attributes: dict, entry_pk: int, node_type: str + cls, + retrieved_attributes: dict, + desired_attributes: set, + entry_pk: int, + node_type: str, ) -> dict: """Build attributes dictionary for OPTIMADE structure resource :param retrieved_attributes: Dict of new attributes, will be updated accordingly :type retrieved_attributes: dict + :param desired_attributes: List of attributes to be built. + :type desired_attributes: set + :param entry_pk: The AiiDA Node's PK :type entry_pk: int + + :param node_type: The AiiDA Node's type + :type node_type: str """ float_fields = { "elements_ratios", @@ -49,22 +58,30 @@ def build_attributes( } # Add existing attributes - missing_attributes = cls.ALL_ATTRIBUTES.copy() existing_attributes = set(retrieved_attributes.keys()) - missing_attributes.difference_update(existing_attributes) + desired_attributes.difference_update(existing_attributes) for field in float_fields: if field in existing_attributes and retrieved_attributes.get(field): retrieved_attributes[field] = hex_to_floats(retrieved_attributes[field]) res = retrieved_attributes.copy() + none_value_attributes = cls.REQUIRED_FIELDS - desired_attributes.union( + existing_attributes + ) + none_value_attributes = { + _ for _ in none_value_attributes if not _.startswith("_") + } + res.update({field: None for field in none_value_attributes}) + # Create and add new attributes - if missing_attributes: + if desired_attributes: translator = cls.TRANSLATORS[node_type](entry_pk) - for attribute in missing_attributes: + + for attribute in desired_attributes: try: create_attribute = getattr(translator, attribute) except AttributeError as exc: - if attribute in cls.REQUIRED_ATTRIBUTES: + if attribute in cls.get_required_fields(): translator = None raise NotImplementedError( f"Parsing required attribute {attribute!r} from " @@ -80,10 +97,11 @@ def build_attributes( ) else: res[attribute] = create_attribute() + # Special post-treatment for `structure_features` all_fields = ( - translator._get_optimade_extras() - ) # pylint: disable=protected-access + translator._get_optimade_extras() # pylint: disable=protected-access + ) all_fields.update(translator.new_attributes) structure_features = all_fields.get("structure_features", []) if all_fields.get("species", None) is None: @@ -97,6 +115,10 @@ def build_attributes( # Some fields were removed translator.new_attributes["structure_features"] = structure_features + translator.new_attributes.update( + {field: None for field in none_value_attributes} + ) + # Store new attributes in `extras` translator.store_attributes() del translator diff --git a/aiida_optimade/models/structures.py b/aiida_optimade/models/structures.py index 2d77a3a4..2342f765 100644 --- a/aiida_optimade/models/structures.py +++ b/aiida_optimade/models/structures.py @@ -1,12 +1,12 @@ # pylint: disable=missing-class-docstring,too-few-public-methods from datetime import datetime - -from pydantic import Field +from typing import Optional from optimade.models import ( StructureResource as OptimadeStructureResource, StructureResourceAttributes as OptimadeStructureResourceAttributes, ) +from optimade.models.utils import OptimadeField, SupportLevel def prefix_provider(string: str) -> str: @@ -21,8 +21,11 @@ def prefix_provider(string: str) -> str: class StructureResourceAttributes(OptimadeStructureResourceAttributes): """Extended StructureResourceAttributes for AiiDA-specific fields""" - ctime: datetime = Field( - ..., description="Creation time of the Node in the AiiDA database." + ctime: Optional[datetime] = OptimadeField( + ..., + description="Creation time of the Node in the AiiDA database.", + support=SupportLevel.SHOULD, + queryable=SupportLevel.MUST, ) class Config: