Coverage for /builds/ase/ase/ase/io/espresso_namelist/namelist.py: 95.31%

1# fmt: off 

2 

3import re 

4import warnings 

5from collections import UserDict 

6from collections.abc import MutableMapping 

7from pathlib import Path 

8 

9from ase.io.espresso_namelist.keys import ALL_KEYS 

10 

11 

12class Namelist(UserDict): 

13 """A case-insensitive dictionary for storing Quantum Espresso namelists. 

14 This class is a subclass of UserDict, which is a wrapper around a regular 

15 dictionary. This allows us to define custom behavior for the dictionary 

16 methods, while still having access to the full dictionary API. 

17 

18 to_string() have been added to handle the conversion of the dictionary 

19 to a string for writing to a file or quick lookup using print(). 

20 

21 to_nested() have been added to convert the dictionary to a nested 

22 dictionary with the correct structure for the specified binary. 

23 """ 

24 def __getitem__(self, key): 

25 return super().__getitem__(key.lower()) 

26 

27 def __setitem__(self, key, value): 

28 super().__setitem__( 

29 key.lower(), Namelist(value) if isinstance( 

30 value, MutableMapping) else value) 

31 

32 def __delitem__(self, key): 

33 super().__delitem__(key.lower()) 

34 

35 @staticmethod 

36 def search_key(to_find, keys): 

37 """Search for a key in the namelist, case-insensitive. 

38 Returns the section and key if found, None otherwise. 

39 """ 

40 for section in keys: 

41 for key in keys[section]: 

42 if re.match(rf"({key})\b(\(+.*\)+)?$", to_find): 

43 return section 

44 

45 def to_string(self, indent=0, list_form=False): 

46 """Format a Namelist object as a string for writing to a file. 

47 Assume sections are ordered (taken care of in namelist construction) 

48 and that repr converts to a QE readable representation (except bools) 

49 

50 Parameters 

51 ---------- 

52 indent : int 

53 Number of spaces to indent each line 

54 list_form : bool 

55 If True, return a list of strings instead of a single string 

56 

57 Returns 

58 ------- 

59 pwi : List[str] | str 

60 Input line for the namelist 

61 """ 

62 pwi = [] 

63 for key, value in self.items(): 

64 if isinstance(value, (Namelist, dict)): 

65 pwi.append(f"{' ' * indent}&{key.upper()}\n") 

66 pwi.extend(Namelist.to_string(value, indent=indent + 3)) 

67 pwi.append(f"{' ' * indent}/\n") 

68 else: 

69 if value is True: 

70 pwi.append(f"{' ' * indent}{key:16} = .true.\n") 

71 elif value is False: 

72 pwi.append(f"{' ' * indent}{key:16} = .false.\n") 

73 elif isinstance(value, Path): 

74 pwi.append(f"{' ' * indent}{key:16} = '{value}'\n") 

75 else: 

76 pwi.append(f"{' ' * indent}{key:16} = {value!r}\n") 

77 if list_form: 

78 return pwi 

79 else: 

80 return "".join(pwi) 

81 

82 def to_nested(self, binary='pw', warn=False, **kwargs): 

83 keys = ALL_KEYS[binary] 

84 

85 constructed_namelist = { 

86 section: self.pop(section, {}) for section in keys 

87 } 

88 

89 constructed_namelist.update({ 

90 key: value for key, value in self.items() 

91 if isinstance(value, Namelist) 

92 }) 

93 

94 unused_keys = [] 

95 for arg_key in list(self): 

96 section = Namelist.search_key(arg_key, keys) 

97 value = self.pop(arg_key) 

98 if section: 

99 constructed_namelist[section][arg_key] = value 

100 else: 

101 unused_keys.append(arg_key) 

102 

103 for arg_key in list(kwargs): 

104 section = Namelist.search_key(arg_key, keys) 

105 value = kwargs.pop(arg_key) 

106 if section: 

107 warnings.warn( 

108 ("Use of kwarg(s) as keyword(s) is deprecated," 

109 "use input_data instead"), 

110 DeprecationWarning, 

111 ) 

112 constructed_namelist[section][arg_key] = value 

113 else: 

114 unused_keys.append(arg_key) 

115 

116 if unused_keys and warn: 

117 warnings.warn( 

118 f"Unused keys: {', '.join(unused_keys)}", 

119 UserWarning, 

120 ) 

121 

122 for section in constructed_namelist: 

123 sorted_section = {} 

124 

125 def sorting_rule(item): 

126 return keys[section].index(item.split('(')[0]) if item.split( 

127 '(')[0] in keys.get(section, {}) else float('inf') 

128 

129 for key in sorted(constructed_namelist[section], key=sorting_rule): 

130 sorted_section[key] = constructed_namelist[section][key] 

131 

132 constructed_namelist[section] = sorted_section 

133 

134 super().update(Namelist(constructed_namelist))