Coverage for /builds/ase/ase/ase/utils/arraywrapper.py: 100.00%

1"""Module for wrapping an array without being an array. 

2 

3This can be desirable because we would like atoms.cell to be like an array, 

4but we don't like atoms.positions @ atoms.cell to also be a cell, and as far 

5as I/we know, it's not possible to subclass array without that kind of thing 

6happening. 

7 

8For most methods and attributes we can just bind the name as a property: 

9 

10 @property 

11 def reshape(self): 

12 return np.asarray(self).reshape 

13 

14For 'in-place' operations like += we need to make sure to return self 

15rather than the array though: 

16 

17 def __iadd__(self, other): 

18 np.asarray(self).__iadd__(other) 

19 return self 

20 

21This module provides the @arraylike decorator which does these things 

22for all the interesting ndarray methods. 

23""" 

24 

25from functools import update_wrapper 

26 

27import numpy as np 

28 

29inplace_methods = [ 

30 '__iadd__', 

31 '__imul__', 

32 '__ipow__', 

33 '__isub__', 

34 '__itruediv__', 

35 '__imatmul__', 

36] 

37 

38forward_methods = [ 

39 '__abs__', 

40 '__add__', 

41 '__contains__', 

42 '__eq__', 

43 '__ge__', 

44 '__getitem__', 

45 '__gt__', 

46 '__hash__', 

47 '__iter__', 

48 '__le__', 

49 '__len__', 

50 '__lt__', 

51 '__mul__', 

52 '__ne__', 

53 '__neg__', 

54 '__pos__', 

55 '__pow__', 

56 '__radd__', 

57 '__rmul__', 

58 '__rpow__', 

59 '__rsub__', 

60 '__rtruediv__', 

61 '__setitem__', 

62 '__sub__', 

63 '__truediv__', 

64] 

65 

66default_methods = [ 

67 '__eq__', 

68 '__le__', 

69 '__lt__', 

70 '__ge__', 

71 '__gt__', 

72 '__ne__', 

73 '__hash__', 

74] 

75 

76if hasattr(np.ndarray, '__matmul__'): 

77 forward_methods += ['__matmul__', '__rmatmul__'] 

78 

79 

80def forward_inplace_call(name): 

81 arraymeth = getattr(np.ndarray, name) 

82 

83 def f(self, obj): 

84 a = self.__array__() 

85 arraymeth(a, obj) 

86 return self 

87 

88 update_wrapper(f, arraymeth) 

89 return f 

90 

91 

92def wrap_array_attribute(name): 

93 wrappee = getattr(np.ndarray, name) 

94 if wrappee is None: # For example, __hash__ is None 

95 assert name == '__hash__' 

96 return None 

97 

98 def attr(self): 

99 array = np.asarray(self) 

100 return getattr(array, name) 

101 

102 update_wrapper(attr, wrappee) 

103 

104 # We don't want to encourage too liberal use of the numpy methods, 

105 # nor do we want the web docs to explode with numpy docstrings or 

106 # break our own doctests. 

107 # 

108 # Therefore we cheat and remove the docstring: 

109 attr.__doc__ = None 

110 return property(attr) 

111 

112 

113def arraylike(cls): 

114 """Decorator for being like an array without being an array. 

115 

116 Poke decorators onto cls so that getting an attribute 

117 really gets that attribute from the wrapped ndarray. 

118 

119 Exceptions are made for in-place methods like +=, *=, etc. 

120 These must return self since otherwise myobj += 1 would 

121 magically turn into an ndarray.""" 

122 for name in inplace_methods: 

123 if hasattr(np.ndarray, name) and not hasattr(cls, name): 

124 meth = forward_inplace_call(name) 

125 setattr(cls, name, meth) 

126 

127 allnames = [name for name in dir(np.ndarray) if not name.startswith('_')] 

128 for name in forward_methods + allnames: 

129 if hasattr(cls, name) and name not in default_methods: 

130 continue # Was overridden -- or there's a conflict. 

131 

132 prop = wrap_array_attribute(name) 

133 setattr(cls, name, prop) 

134 return cls