Coverage for /builds/debichem-team/python-ase/ase/io/espresso_namelist/namelist.py: 95.31%

1import re 

2import warnings 

3from collections import UserDict 

4from collections.abc import MutableMapping 

5from pathlib import Path 

6 

7from ase.io.espresso_namelist.keys import ALL_KEYS 

8 

9 

10class Namelist(UserDict): 

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

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

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

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

15 

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

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

18 

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

20 dictionary with the correct structure for the specified binary. 

21 """ 

22 def __getitem__(self, key): 

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

24 

25 def __setitem__(self, key, value): 

26 super().__setitem__( 

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

28 value, MutableMapping) else value) 

29 

30 def __delitem__(self, key): 

31 super().__delitem__(key.lower()) 

32 

33 @staticmethod 

34 def search_key(to_find, keys): 

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

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

37 """ 

38 for section in keys: 

39 for key in keys[section]: 

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

41 return section 

42 

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

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

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

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

47 

48 Parameters 

49 ---------- 

50 indent : int 

51 Number of spaces to indent each line 

52 list_form : bool 

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

54 

55 Returns 

56 ------- 

57 pwi : List[str] | str 

58 Input line for the namelist 

59 """ 

60 pwi = [] 

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

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

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

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

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

66 else: 

67 if value is True: 

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

69 elif value is False: 

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

71 elif isinstance(value, Path): 

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

73 else: 

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

75 if list_form: 

76 return pwi 

77 else: 

78 return "".join(pwi) 

79 

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

81 keys = ALL_KEYS[binary] 

82 

83 constructed_namelist = { 

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

85 } 

86 

87 constructed_namelist.update({ 

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

89 if isinstance(value, Namelist) 

90 }) 

91 

92 unused_keys = [] 

93 for arg_key in list(self): 

94 section = Namelist.search_key(arg_key, keys) 

95 value = self.pop(arg_key) 

96 if section: 

97 constructed_namelist[section][arg_key] = value 

98 else: 

99 unused_keys.append(arg_key) 

100 

101 for arg_key in list(kwargs): 

102 section = Namelist.search_key(arg_key, keys) 

103 value = kwargs.pop(arg_key) 

104 if section: 

105 warnings.warn( 

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

107 "use input_data instead"), 

108 DeprecationWarning, 

109 ) 

110 constructed_namelist[section][arg_key] = value 

111 else: 

112 unused_keys.append(arg_key) 

113 

114 if unused_keys and warn: 

115 warnings.warn( 

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

117 UserWarning, 

118 ) 

119 

120 for section in constructed_namelist: 

121 sorted_section = {} 

122 

123 def sorting_rule(item): 

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

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

126 

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

128 sorted_section[key] = constructed_namelist[section][key] 

129 

130 constructed_namelist[section] = sorted_section 

131 

132 super().update(Namelist(constructed_namelist))