diff --git a/src/aiida_pseudo/groups/family/pseudo.py b/src/aiida_pseudo/groups/family/pseudo.py index d2e4a6a..b6365e1 100644 --- a/src/aiida_pseudo/groups/family/pseudo.py +++ b/src/aiida_pseudo/groups/family/pseudo.py @@ -315,8 +315,8 @@ def get_pseudos( self, *, elements: Optional[Union[List[str], Tuple[str]]] = None, - structure: Union[structures_classes] = None, - ) -> Mapping[str, Union[structures_classes]]: + structure: Union[structures_classes] = None, # type: ignore + ) -> Mapping[str, Union[structures_classes]]: # type: ignore """Return the mapping of kind names on pseudo potential data nodes for the given list of elements or structure. :param elements: list of element symbols. @@ -337,6 +337,11 @@ def get_pseudos( raise ValueError(f'structure is of type {type(structure)} but should be of: {structures_classes}') if structure is not None: - return {kind.name: self.get_pseudo(kind.symbol) for kind in structure.kinds} + if isinstance(structure, LegacyStructureData): + return {kind.name: self.get_pseudo(kind.symbol) for kind in structure.kinds} + elif structure.properties.kind_names: + return {kind.kind_name: self.get_pseudo(kind.symbol) for kind in structure.kinds} + else: + return {symbol: self.get_pseudo(symbol) for symbol in structure.symbols} return {element: self.get_pseudo(element) for element in elements} diff --git a/src/aiida_pseudo/groups/mixins/cutoffs.py b/src/aiida_pseudo/groups/mixins/cutoffs.py index 3f37f57..14fe418 100644 --- a/src/aiida_pseudo/groups/mixins/cutoffs.py +++ b/src/aiida_pseudo/groups/mixins/cutoffs.py @@ -2,7 +2,7 @@ import warnings from typing import Optional -from aiida.common.exceptions import MissingEntryPointError +from aiida.common import exceptions from aiida.common.lang import type_check from aiida.plugins import DataFactory @@ -12,7 +12,7 @@ try: StructureData = DataFactory('atomistic.structure') -except MissingEntryPointError: +except exceptions.MissingEntryPointError: structures_classes = (LegacyStructureData,) else: structures_classes = (LegacyStructureData, StructureData) @@ -286,7 +286,7 @@ def get_recommended_cutoffs(self, *, elements=None, structure=None, stringency=N raise ValueError('at least one and only one of `elements` or `structure` should be defined') type_check(elements, (tuple, str), allow_none=True) - type_check(structure, (structures_classes), allow_none=True) + type_check(structure, structures_classes, allow_none=True) if unit is not None: self.validate_cutoffs_unit(unit)