Coverage for /builds/hweiske/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 def __getitem__(self, key): 

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

13 

14 def __setitem__(self, key, value): 

15 super().__setitem__( 

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

17 value, MutableMapping) else value) 

18 

19 def __delitem__(self, key): 

20 super().__delitem__(key.lower()) 

21 

22 @staticmethod 

23 def search_key(to_find, keys): 

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

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

26 """ 

27 for section in keys: 

28 for key in keys[section]: 

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

30 return section 

31 

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

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

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

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

36 

37 Parameters 

38 ---------- 

39 indent : int 

40 Number of spaces to indent each line 

41 list_form : bool 

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

43 

44 Returns 

45 ------- 

46 pwi : List[str] | str 

47 Input line for the namelist 

48 """ 

49 pwi = [] 

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

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

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

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

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

55 else: 

56 if value is True: 

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

58 elif value is False: 

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

60 elif isinstance(value, Path): 

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

62 else: 

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

64 if list_form: 

65 return pwi 

66 else: 

67 return "".join(pwi) 

68 

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

70 keys = ALL_KEYS[binary] 

71 

72 constructed_namelist = { 

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

74 } 

75 

76 constructed_namelist.update({ 

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

78 if isinstance(value, Namelist) 

79 }) 

80 

81 unused_keys = [] 

82 for arg_key in list(self): 

83 section = Namelist.search_key(arg_key, keys) 

84 value = self.pop(arg_key) 

85 if section: 

86 constructed_namelist[section][arg_key] = value 

87 else: 

88 unused_keys.append(arg_key) 

89 

90 for arg_key in list(kwargs): 

91 section = Namelist.search_key(arg_key, keys) 

92 value = kwargs.pop(arg_key) 

93 if section: 

94 warnings.warn( 

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

96 "use input_data instead"), 

97 DeprecationWarning, 

98 ) 

99 constructed_namelist[section][arg_key] = value 

100 else: 

101 unused_keys.append(arg_key) 

102 

103 if unused_keys and warn: 

104 warnings.warn( 

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

106 UserWarning, 

107 ) 

108 

109 for section in constructed_namelist: 

110 sorted_section = {} 

111 

112 def sorting_rule(item): 

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

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

115 

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

117 sorted_section[key] = constructed_namelist[section][key] 

118 

119 constructed_namelist[section] = sorted_section 

120 

121 super().update(Namelist(constructed_namelist))