Coverage for /usr/lib/python3/dist-packages/sympy/core/traversal.py: 49%

1from .basic import Basic 

2from .sorting import ordered 

3from .sympify import sympify 

4from sympy.utilities.iterables import iterable 

5 

6 

7 

8def iterargs(expr): 

9 """Yield the args of a Basic object in a breadth-first traversal. 

10 Depth-traversal stops if `arg.args` is either empty or is not 

11 an iterable. 

12 

13 Examples 

14 ======== 

15 

16 >>> from sympy import Integral, Function 

17 >>> from sympy.abc import x 

18 >>> f = Function('f') 

19 >>> from sympy.core.traversal import iterargs 

20 >>> list(iterargs(Integral(f(x), (f(x), 1)))) 

21 [Integral(f(x), (f(x), 1)), f(x), (f(x), 1), x, f(x), 1, x] 

22 

23 See Also 

24 ======== 

25 iterfreeargs, preorder_traversal 

26 """ 

27 args = [expr] 

28 for i in args: 

29 yield i 

30 try: 

31 args.extend(i.args) 

32 except TypeError: 

33 pass # for cases like f being an arg 

34 

35 

36def iterfreeargs(expr, _first=True): 

37 """Yield the args of a Basic object in a breadth-first traversal. 

38 Depth-traversal stops if `arg.args` is either empty or is not 

39 an iterable. The bound objects of an expression will be returned 

40 as canonical variables. 

41 

42 Examples 

43 ======== 

44 

45 >>> from sympy import Integral, Function 

46 >>> from sympy.abc import x 

47 >>> f = Function('f') 

48 >>> from sympy.core.traversal import iterfreeargs 

49 >>> list(iterfreeargs(Integral(f(x), (f(x), 1)))) 

50 [Integral(f(x), (f(x), 1)), 1] 

51 

52 See Also 

53 ======== 

54 iterargs, preorder_traversal 

55 """ 

56 args = [expr] 

57 for i in args: 

58 yield i 

59 if _first and hasattr(i, 'bound_symbols'): 

60 void = i.canonical_variables.values() 

61 for i in iterfreeargs(i.as_dummy(), _first=False): 

62 if not i.has(*void): 

63 yield i 

64 try: 

65 args.extend(i.args) 

66 except TypeError: 

67 pass # for cases like f being an arg 

68 

69 

70class preorder_traversal: 

71 """ 

72 Do a pre-order traversal of a tree. 

73 

74 This iterator recursively yields nodes that it has visited in a pre-order 

75 fashion. That is, it yields the current node then descends through the 

76 tree breadth-first to yield all of a node's children's pre-order 

77 traversal. 

78 

79 

80 For an expression, the order of the traversal depends on the order of 

81 .args, which in many cases can be arbitrary. 

82 

83 Parameters 

84 ========== 

85 node : SymPy expression 

86 The expression to traverse. 

87 keys : (default None) sort key(s) 

88 The key(s) used to sort args of Basic objects. When None, args of Basic 

89 objects are processed in arbitrary order. If key is defined, it will 

90 be passed along to ordered() as the only key(s) to use to sort the 

91 arguments; if ``key`` is simply True then the default keys of ordered 

92 will be used. 

93 

94 Yields 

95 ====== 

96 subtree : SymPy expression 

97 All of the subtrees in the tree. 

98 

99 Examples 

100 ======== 

101 

102 >>> from sympy import preorder_traversal, symbols 

103 >>> x, y, z = symbols('x y z') 

104 

105 The nodes are returned in the order that they are encountered unless key 

106 is given; simply passing key=True will guarantee that the traversal is 

107 unique. 

108 

109 >>> list(preorder_traversal((x + y)*z, keys=None)) # doctest: +SKIP 

110 [z*(x + y), z, x + y, y, x] 

111 >>> list(preorder_traversal((x + y)*z, keys=True)) 

112 [z*(x + y), z, x + y, x, y] 

113 

114 """ 

115 def __init__(self, node, keys=None): 

116 self._skip_flag = False 

117 self._pt = self._preorder_traversal(node, keys) 

118 

119 def _preorder_traversal(self, node, keys): 

120 yield node 

121 if self._skip_flag: 

122 self._skip_flag = False 

123 return 

124 if isinstance(node, Basic): 

125 if not keys and hasattr(node, '_argset'): 

126 # LatticeOp keeps args as a set. We should use this if we 

127 # don't care about the order, to prevent unnecessary sorting. 

128 args = node._argset 

129 else: 

130 args = node.args 

131 if keys: 

132 if keys != True: 

133 args = ordered(args, keys, default=False) 

134 else: 

135 args = ordered(args) 

136 for arg in args: 

137 yield from self._preorder_traversal(arg, keys) 

138 elif iterable(node): 

139 for item in node: 

140 yield from self._preorder_traversal(item, keys) 

141 

142 def skip(self): 

143 """ 

144 Skip yielding current node's (last yielded node's) subtrees. 

145 

146 Examples 

147 ======== 

148 

149 >>> from sympy import preorder_traversal, symbols 

150 >>> x, y, z = symbols('x y z') 

151 >>> pt = preorder_traversal((x + y*z)*z) 

152 >>> for i in pt: 

153 ... print(i) 

154 ... if i == x + y*z: 

155 ... pt.skip() 

156 z*(x + y*z) 

157 z 

158 x + y*z 

159 """ 

160 self._skip_flag = True 

161 

162 def __next__(self): 

163 return next(self._pt) 

164 

165 def __iter__(self): 

166 return self 

167 

168 

169def use(expr, func, level=0, args=(), kwargs={}): 

170 """ 

171 Use ``func`` to transform ``expr`` at the given level. 

172 

173 Examples 

174 ======== 

175 

176 >>> from sympy import use, expand 

177 >>> from sympy.abc import x, y 

178 

179 >>> f = (x + y)**2*x + 1 

180 

181 >>> use(f, expand, level=2) 

182 x*(x**2 + 2*x*y + y**2) + 1 

183 >>> expand(f) 

184 x**3 + 2*x**2*y + x*y**2 + 1 

185 

186 """ 

187 def _use(expr, level): 

188 if not level: 

189 return func(expr, *args, **kwargs) 

190 else: 

191 if expr.is_Atom: 

192 return expr 

193 else: 

194 level -= 1 

195 _args = [_use(arg, level) for arg in expr.args] 

196 return expr.__class__(*_args) 

197 

198 return _use(sympify(expr), level) 

199 

200 

201def walk(e, *target): 

202 """Iterate through the args that are the given types (target) and 

203 return a list of the args that were traversed; arguments 

204 that are not of the specified types are not traversed. 

205 

206 Examples 

207 ======== 

208 

209 >>> from sympy.core.traversal import walk 

210 >>> from sympy import Min, Max 

211 >>> from sympy.abc import x, y, z 

212 >>> list(walk(Min(x, Max(y, Min(1, z))), Min)) 

213 [Min(x, Max(y, Min(1, z)))] 

214 >>> list(walk(Min(x, Max(y, Min(1, z))), Min, Max)) 

215 [Min(x, Max(y, Min(1, z))), Max(y, Min(1, z)), Min(1, z)] 

216 

217 See Also 

218 ======== 

219 

220 bottom_up 

221 """ 

222 if isinstance(e, target): 

223 yield e 

224 for i in e.args: 

225 yield from walk(i, *target) 

226 

227 

228def bottom_up(rv, F, atoms=False, nonbasic=False): 

229 """Apply ``F`` to all expressions in an expression tree from the 

230 bottom up. If ``atoms`` is True, apply ``F`` even if there are no args; 

231 if ``nonbasic`` is True, try to apply ``F`` to non-Basic objects. 

232 """ 

233 args = getattr(rv, 'args', None) 

234 if args is not None: 

235 if args: 

236 args = tuple([bottom_up(a, F, atoms, nonbasic) for a in args]) 

237 if args != rv.args: 

238 rv = rv.func(*args) 

239 rv = F(rv) 

240 elif atoms: 

241 rv = F(rv) 

242 else: 

243 if nonbasic: 

244 try: 

245 rv = F(rv) 

246 except TypeError: 

247 pass 

248 

249 return rv 

250 

251 

252def postorder_traversal(node, keys=None): 

253 """ 

254 Do a postorder traversal of a tree. 

255 

256 This generator recursively yields nodes that it has visited in a postorder 

257 fashion. That is, it descends through the tree depth-first to yield all of 

258 a node's children's postorder traversal before yielding the node itself. 

259 

260 Parameters 

261 ========== 

262 

263 node : SymPy expression 

264 The expression to traverse. 

265 keys : (default None) sort key(s) 

266 The key(s) used to sort args of Basic objects. When None, args of Basic 

267 objects are processed in arbitrary order. If key is defined, it will 

268 be passed along to ordered() as the only key(s) to use to sort the 

269 arguments; if ``key`` is simply True then the default keys of 

270 ``ordered`` will be used (node count and default_sort_key). 

271 

272 Yields 

273 ====== 

274 subtree : SymPy expression 

275 All of the subtrees in the tree. 

276 

277 Examples 

278 ======== 

279 

280 >>> from sympy import postorder_traversal 

281 >>> from sympy.abc import w, x, y, z 

282 

283 The nodes are returned in the order that they are encountered unless key 

284 is given; simply passing key=True will guarantee that the traversal is 

285 unique. 

286 

287 >>> list(postorder_traversal(w + (x + y)*z)) # doctest: +SKIP 

288 [z, y, x, x + y, z*(x + y), w, w + z*(x + y)] 

289 >>> list(postorder_traversal(w + (x + y)*z, keys=True)) 

290 [w, z, x, y, x + y, z*(x + y), w + z*(x + y)] 

291 

292 

293 """ 

294 if isinstance(node, Basic): 

295 args = node.args 

296 if keys: 

297 if keys != True: 

298 args = ordered(args, keys, default=False) 

299 else: 

300 args = ordered(args) 

301 for arg in args: 

302 yield from postorder_traversal(arg, keys) 

303 elif iterable(node): 

304 for item in node: 

305 yield from postorder_traversal(item, keys) 

306 yield node