test_ast.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661
  1. import math
  2. from sympy.core.containers import Tuple
  3. from sympy.core.numbers import nan, oo, Float, Integer
  4. from sympy.core.relational import Lt
  5. from sympy.core.symbol import symbols, Symbol
  6. from sympy.functions.elementary.trigonometric import sin
  7. from sympy.matrices.dense import Matrix
  8. from sympy.matrices.expressions.matexpr import MatrixSymbol
  9. from sympy.sets.fancysets import Range
  10. from sympy.tensor.indexed import Idx, IndexedBase
  11. from sympy.testing.pytest import raises
  12. from sympy.codegen.ast import (
  13. Assignment, Attribute, aug_assign, CodeBlock, For, Type, Variable, Pointer, Declaration,
  14. AddAugmentedAssignment, SubAugmentedAssignment, MulAugmentedAssignment,
  15. DivAugmentedAssignment, ModAugmentedAssignment, value_const, pointer_const,
  16. integer, real, complex_, int8, uint8, float16 as f16, float32 as f32,
  17. float64 as f64, float80 as f80, float128 as f128, complex64 as c64, complex128 as c128,
  18. While, Scope, String, Print, QuotedString, FunctionPrototype, FunctionDefinition, Return,
  19. FunctionCall, untyped, IntBaseType, intc, Node, none, NoneToken, Token, Comment
  20. )
  21. x, y, z, t, x0, x1, x2, a, b = symbols("x, y, z, t, x0, x1, x2, a, b")
  22. n = symbols("n", integer=True)
  23. A = MatrixSymbol('A', 3, 1)
  24. mat = Matrix([1, 2, 3])
  25. B = IndexedBase('B')
  26. i = Idx("i", n)
  27. A22 = MatrixSymbol('A22',2,2)
  28. B22 = MatrixSymbol('B22',2,2)
  29. def test_Assignment():
  30. # Here we just do things to show they don't error
  31. Assignment(x, y)
  32. Assignment(x, 0)
  33. Assignment(A, mat)
  34. Assignment(A[1,0], 0)
  35. Assignment(A[1,0], x)
  36. Assignment(B[i], x)
  37. Assignment(B[i], 0)
  38. a = Assignment(x, y)
  39. assert a.func(*a.args) == a
  40. assert a.op == ':='
  41. # Here we test things to show that they error
  42. # Matrix to scalar
  43. raises(ValueError, lambda: Assignment(B[i], A))
  44. raises(ValueError, lambda: Assignment(B[i], mat))
  45. raises(ValueError, lambda: Assignment(x, mat))
  46. raises(ValueError, lambda: Assignment(x, A))
  47. raises(ValueError, lambda: Assignment(A[1,0], mat))
  48. # Scalar to matrix
  49. raises(ValueError, lambda: Assignment(A, x))
  50. raises(ValueError, lambda: Assignment(A, 0))
  51. # Non-atomic lhs
  52. raises(TypeError, lambda: Assignment(mat, A))
  53. raises(TypeError, lambda: Assignment(0, x))
  54. raises(TypeError, lambda: Assignment(x*x, 1))
  55. raises(TypeError, lambda: Assignment(A + A, mat))
  56. raises(TypeError, lambda: Assignment(B, 0))
  57. def test_AugAssign():
  58. # Here we just do things to show they don't error
  59. aug_assign(x, '+', y)
  60. aug_assign(x, '+', 0)
  61. aug_assign(A, '+', mat)
  62. aug_assign(A[1, 0], '+', 0)
  63. aug_assign(A[1, 0], '+', x)
  64. aug_assign(B[i], '+', x)
  65. aug_assign(B[i], '+', 0)
  66. # Check creation via aug_assign vs constructor
  67. for binop, cls in [
  68. ('+', AddAugmentedAssignment),
  69. ('-', SubAugmentedAssignment),
  70. ('*', MulAugmentedAssignment),
  71. ('/', DivAugmentedAssignment),
  72. ('%', ModAugmentedAssignment),
  73. ]:
  74. a = aug_assign(x, binop, y)
  75. b = cls(x, y)
  76. assert a.func(*a.args) == a == b
  77. assert a.binop == binop
  78. assert a.op == binop + '='
  79. # Here we test things to show that they error
  80. # Matrix to scalar
  81. raises(ValueError, lambda: aug_assign(B[i], '+', A))
  82. raises(ValueError, lambda: aug_assign(B[i], '+', mat))
  83. raises(ValueError, lambda: aug_assign(x, '+', mat))
  84. raises(ValueError, lambda: aug_assign(x, '+', A))
  85. raises(ValueError, lambda: aug_assign(A[1, 0], '+', mat))
  86. # Scalar to matrix
  87. raises(ValueError, lambda: aug_assign(A, '+', x))
  88. raises(ValueError, lambda: aug_assign(A, '+', 0))
  89. # Non-atomic lhs
  90. raises(TypeError, lambda: aug_assign(mat, '+', A))
  91. raises(TypeError, lambda: aug_assign(0, '+', x))
  92. raises(TypeError, lambda: aug_assign(x * x, '+', 1))
  93. raises(TypeError, lambda: aug_assign(A + A, '+', mat))
  94. raises(TypeError, lambda: aug_assign(B, '+', 0))
  95. def test_Assignment_printing():
  96. assignment_classes = [
  97. Assignment,
  98. AddAugmentedAssignment,
  99. SubAugmentedAssignment,
  100. MulAugmentedAssignment,
  101. DivAugmentedAssignment,
  102. ModAugmentedAssignment,
  103. ]
  104. pairs = [
  105. (x, 2 * y + 2),
  106. (B[i], x),
  107. (A22, B22),
  108. (A[0, 0], x),
  109. ]
  110. for cls in assignment_classes:
  111. for lhs, rhs in pairs:
  112. a = cls(lhs, rhs)
  113. assert repr(a) == '%s(%s, %s)' % (cls.__name__, repr(lhs), repr(rhs))
  114. def test_CodeBlock():
  115. c = CodeBlock(Assignment(x, 1), Assignment(y, x + 1))
  116. assert c.func(*c.args) == c
  117. assert c.left_hand_sides == Tuple(x, y)
  118. assert c.right_hand_sides == Tuple(1, x + 1)
  119. def test_CodeBlock_topological_sort():
  120. assignments = [
  121. Assignment(x, y + z),
  122. Assignment(z, 1),
  123. Assignment(t, x),
  124. Assignment(y, 2),
  125. ]
  126. ordered_assignments = [
  127. # Note that the unrelated z=1 and y=2 are kept in that order
  128. Assignment(z, 1),
  129. Assignment(y, 2),
  130. Assignment(x, y + z),
  131. Assignment(t, x),
  132. ]
  133. c1 = CodeBlock.topological_sort(assignments)
  134. assert c1 == CodeBlock(*ordered_assignments)
  135. # Cycle
  136. invalid_assignments = [
  137. Assignment(x, y + z),
  138. Assignment(z, 1),
  139. Assignment(y, x),
  140. Assignment(y, 2),
  141. ]
  142. raises(ValueError, lambda: CodeBlock.topological_sort(invalid_assignments))
  143. # Free symbols
  144. free_assignments = [
  145. Assignment(x, y + z),
  146. Assignment(z, a * b),
  147. Assignment(t, x),
  148. Assignment(y, b + 3),
  149. ]
  150. free_assignments_ordered = [
  151. Assignment(z, a * b),
  152. Assignment(y, b + 3),
  153. Assignment(x, y + z),
  154. Assignment(t, x),
  155. ]
  156. c2 = CodeBlock.topological_sort(free_assignments)
  157. assert c2 == CodeBlock(*free_assignments_ordered)
  158. def test_CodeBlock_free_symbols():
  159. c1 = CodeBlock(
  160. Assignment(x, y + z),
  161. Assignment(z, 1),
  162. Assignment(t, x),
  163. Assignment(y, 2),
  164. )
  165. assert c1.free_symbols == set()
  166. c2 = CodeBlock(
  167. Assignment(x, y + z),
  168. Assignment(z, a * b),
  169. Assignment(t, x),
  170. Assignment(y, b + 3),
  171. )
  172. assert c2.free_symbols == {a, b}
  173. def test_CodeBlock_cse():
  174. c1 = CodeBlock(
  175. Assignment(y, 1),
  176. Assignment(x, sin(y)),
  177. Assignment(z, sin(y)),
  178. Assignment(t, x*z),
  179. )
  180. assert c1.cse() == CodeBlock(
  181. Assignment(y, 1),
  182. Assignment(x0, sin(y)),
  183. Assignment(x, x0),
  184. Assignment(z, x0),
  185. Assignment(t, x*z),
  186. )
  187. # Multiple assignments to same symbol not supported
  188. raises(NotImplementedError, lambda: CodeBlock(
  189. Assignment(x, 1),
  190. Assignment(y, 1), Assignment(y, 2)
  191. ).cse())
  192. # Check auto-generated symbols do not collide with existing ones
  193. c2 = CodeBlock(
  194. Assignment(x0, sin(y) + 1),
  195. Assignment(x1, 2 * sin(y)),
  196. Assignment(z, x * y),
  197. )
  198. assert c2.cse() == CodeBlock(
  199. Assignment(x2, sin(y)),
  200. Assignment(x0, x2 + 1),
  201. Assignment(x1, 2 * x2),
  202. Assignment(z, x * y),
  203. )
  204. def test_CodeBlock_cse__issue_14118():
  205. # see https://github.com/sympy/sympy/issues/14118
  206. c = CodeBlock(
  207. Assignment(A22, Matrix([[x, sin(y)],[3, 4]])),
  208. Assignment(B22, Matrix([[sin(y), 2*sin(y)], [sin(y)**2, 7]]))
  209. )
  210. assert c.cse() == CodeBlock(
  211. Assignment(x0, sin(y)),
  212. Assignment(A22, Matrix([[x, x0],[3, 4]])),
  213. Assignment(B22, Matrix([[x0, 2*x0], [x0**2, 7]]))
  214. )
  215. def test_For():
  216. f = For(n, Range(0, 3), (Assignment(A[n, 0], x + n), aug_assign(x, '+', y)))
  217. f = For(n, (1, 2, 3, 4, 5), (Assignment(A[n, 0], x + n),))
  218. assert f.func(*f.args) == f
  219. raises(TypeError, lambda: For(n, x, (x + y,)))
  220. def test_none():
  221. assert none.is_Atom
  222. assert none == none
  223. class Foo(Token):
  224. pass
  225. foo = Foo()
  226. assert foo != none
  227. assert none == None
  228. assert none == NoneToken()
  229. assert none.func(*none.args) == none
  230. def test_String():
  231. st = String('foobar')
  232. assert st.is_Atom
  233. assert st == String('foobar')
  234. assert st.text == 'foobar'
  235. assert st.func(**st.kwargs()) == st
  236. assert st.func(*st.args) == st
  237. class Signifier(String):
  238. pass
  239. si = Signifier('foobar')
  240. assert si != st
  241. assert si.text == st.text
  242. s = String('foo')
  243. assert str(s) == 'foo'
  244. assert repr(s) == "String('foo')"
  245. def test_Comment():
  246. c = Comment('foobar')
  247. assert c.text == 'foobar'
  248. assert str(c) == 'foobar'
  249. def test_Node():
  250. n = Node()
  251. assert n == Node()
  252. assert n.func(*n.args) == n
  253. def test_Type():
  254. t = Type('MyType')
  255. assert len(t.args) == 1
  256. assert t.name == String('MyType')
  257. assert str(t) == 'MyType'
  258. assert repr(t) == "Type(String('MyType'))"
  259. assert Type(t) == t
  260. assert t.func(*t.args) == t
  261. t1 = Type('t1')
  262. t2 = Type('t2')
  263. assert t1 != t2
  264. assert t1 == t1 and t2 == t2
  265. t1b = Type('t1')
  266. assert t1 == t1b
  267. assert t2 != t1b
  268. def test_Type__from_expr():
  269. assert Type.from_expr(i) == integer
  270. u = symbols('u', real=True)
  271. assert Type.from_expr(u) == real
  272. assert Type.from_expr(n) == integer
  273. assert Type.from_expr(3) == integer
  274. assert Type.from_expr(3.0) == real
  275. assert Type.from_expr(3+1j) == complex_
  276. raises(ValueError, lambda: Type.from_expr(sum))
  277. def test_Type__cast_check__integers():
  278. # Rounding
  279. raises(ValueError, lambda: integer.cast_check(3.5))
  280. assert integer.cast_check('3') == 3
  281. assert integer.cast_check(Float('3.0000000000000000000')) == 3
  282. assert integer.cast_check(Float('3.0000000000000000001')) == 3 # unintuitive maybe?
  283. # Range
  284. assert int8.cast_check(127.0) == 127
  285. raises(ValueError, lambda: int8.cast_check(128))
  286. assert int8.cast_check(-128) == -128
  287. raises(ValueError, lambda: int8.cast_check(-129))
  288. assert uint8.cast_check(0) == 0
  289. assert uint8.cast_check(128) == 128
  290. raises(ValueError, lambda: uint8.cast_check(256.0))
  291. raises(ValueError, lambda: uint8.cast_check(-1))
  292. def test_Attribute():
  293. noexcept = Attribute('noexcept')
  294. assert noexcept == Attribute('noexcept')
  295. alignas16 = Attribute('alignas', [16])
  296. alignas32 = Attribute('alignas', [32])
  297. assert alignas16 != alignas32
  298. assert alignas16.func(*alignas16.args) == alignas16
  299. def test_Variable():
  300. v = Variable(x, type=real)
  301. assert v == Variable(v)
  302. assert v == Variable('x', type=real)
  303. assert v.symbol == x
  304. assert v.type == real
  305. assert value_const not in v.attrs
  306. assert v.func(*v.args) == v
  307. assert str(v) == 'Variable(x, type=real)'
  308. w = Variable(y, f32, attrs={value_const})
  309. assert w.symbol == y
  310. assert w.type == f32
  311. assert value_const in w.attrs
  312. assert w.func(*w.args) == w
  313. v_n = Variable(n, type=Type.from_expr(n))
  314. assert v_n.type == integer
  315. assert v_n.func(*v_n.args) == v_n
  316. v_i = Variable(i, type=Type.from_expr(n))
  317. assert v_i.type == integer
  318. assert v_i != v_n
  319. a_i = Variable.deduced(i)
  320. assert a_i.type == integer
  321. assert Variable.deduced(Symbol('x', real=True)).type == real
  322. assert a_i.func(*a_i.args) == a_i
  323. v_n2 = Variable.deduced(n, value=3.5, cast_check=False)
  324. assert v_n2.func(*v_n2.args) == v_n2
  325. assert abs(v_n2.value - 3.5) < 1e-15
  326. raises(ValueError, lambda: Variable.deduced(n, value=3.5, cast_check=True))
  327. v_n3 = Variable.deduced(n)
  328. assert v_n3.type == integer
  329. assert str(v_n3) == 'Variable(n, type=integer)'
  330. assert Variable.deduced(z, value=3).type == integer
  331. assert Variable.deduced(z, value=3.0).type == real
  332. assert Variable.deduced(z, value=3.0+1j).type == complex_
  333. def test_Pointer():
  334. p = Pointer(x)
  335. assert p.symbol == x
  336. assert p.type == untyped
  337. assert value_const not in p.attrs
  338. assert pointer_const not in p.attrs
  339. assert p.func(*p.args) == p
  340. u = symbols('u', real=True)
  341. pu = Pointer(u, type=Type.from_expr(u), attrs={value_const, pointer_const})
  342. assert pu.symbol is u
  343. assert pu.type == real
  344. assert value_const in pu.attrs
  345. assert pointer_const in pu.attrs
  346. assert pu.func(*pu.args) == pu
  347. i = symbols('i', integer=True)
  348. deref = pu[i]
  349. assert deref.indices == (i,)
  350. def test_Declaration():
  351. u = symbols('u', real=True)
  352. vu = Variable(u, type=Type.from_expr(u))
  353. assert Declaration(vu).variable.type == real
  354. vn = Variable(n, type=Type.from_expr(n))
  355. assert Declaration(vn).variable.type == integer
  356. # PR 19107, does not allow comparison between expressions and Basic
  357. # lt = StrictLessThan(vu, vn)
  358. # assert isinstance(lt, StrictLessThan)
  359. vuc = Variable(u, Type.from_expr(u), value=3.0, attrs={value_const})
  360. assert value_const in vuc.attrs
  361. assert pointer_const not in vuc.attrs
  362. decl = Declaration(vuc)
  363. assert decl.variable == vuc
  364. assert isinstance(decl.variable.value, Float)
  365. assert decl.variable.value == 3.0
  366. assert decl.func(*decl.args) == decl
  367. assert vuc.as_Declaration() == decl
  368. assert vuc.as_Declaration(value=None, attrs=None) == Declaration(vu)
  369. vy = Variable(y, type=integer, value=3)
  370. decl2 = Declaration(vy)
  371. assert decl2.variable == vy
  372. assert decl2.variable.value == Integer(3)
  373. vi = Variable(i, type=Type.from_expr(i), value=3.0)
  374. decl3 = Declaration(vi)
  375. assert decl3.variable.type == integer
  376. assert decl3.variable.value == 3.0
  377. raises(ValueError, lambda: Declaration(vi, 42))
  378. def test_IntBaseType():
  379. assert intc.name == String('intc')
  380. assert intc.args == (intc.name,)
  381. assert str(IntBaseType('a').name) == 'a'
  382. def test_FloatType():
  383. assert f16.dig == 3
  384. assert f32.dig == 6
  385. assert f64.dig == 15
  386. assert f80.dig == 18
  387. assert f128.dig == 33
  388. assert f16.decimal_dig == 5
  389. assert f32.decimal_dig == 9
  390. assert f64.decimal_dig == 17
  391. assert f80.decimal_dig == 21
  392. assert f128.decimal_dig == 36
  393. assert f16.max_exponent == 16
  394. assert f32.max_exponent == 128
  395. assert f64.max_exponent == 1024
  396. assert f80.max_exponent == 16384
  397. assert f128.max_exponent == 16384
  398. assert f16.min_exponent == -13
  399. assert f32.min_exponent == -125
  400. assert f64.min_exponent == -1021
  401. assert f80.min_exponent == -16381
  402. assert f128.min_exponent == -16381
  403. assert abs(f16.eps / Float('0.00097656', precision=16) - 1) < 0.1*10**-f16.dig
  404. assert abs(f32.eps / Float('1.1920929e-07', precision=32) - 1) < 0.1*10**-f32.dig
  405. assert abs(f64.eps / Float('2.2204460492503131e-16', precision=64) - 1) < 0.1*10**-f64.dig
  406. assert abs(f80.eps / Float('1.08420217248550443401e-19', precision=80) - 1) < 0.1*10**-f80.dig
  407. assert abs(f128.eps / Float(' 1.92592994438723585305597794258492732e-34', precision=128) - 1) < 0.1*10**-f128.dig
  408. assert abs(f16.max / Float('65504', precision=16) - 1) < .1*10**-f16.dig
  409. assert abs(f32.max / Float('3.40282347e+38', precision=32) - 1) < 0.1*10**-f32.dig
  410. assert abs(f64.max / Float('1.79769313486231571e+308', precision=64) - 1) < 0.1*10**-f64.dig # cf. np.finfo(np.float64).max
  411. assert abs(f80.max / Float('1.18973149535723176502e+4932', precision=80) - 1) < 0.1*10**-f80.dig
  412. assert abs(f128.max / Float('1.18973149535723176508575932662800702e+4932', precision=128) - 1) < 0.1*10**-f128.dig
  413. # cf. np.finfo(np.float32).tiny
  414. assert abs(f16.tiny / Float('6.1035e-05', precision=16) - 1) < 0.1*10**-f16.dig
  415. assert abs(f32.tiny / Float('1.17549435e-38', precision=32) - 1) < 0.1*10**-f32.dig
  416. assert abs(f64.tiny / Float('2.22507385850720138e-308', precision=64) - 1) < 0.1*10**-f64.dig
  417. assert abs(f80.tiny / Float('3.36210314311209350626e-4932', precision=80) - 1) < 0.1*10**-f80.dig
  418. assert abs(f128.tiny / Float('3.3621031431120935062626778173217526e-4932', precision=128) - 1) < 0.1*10**-f128.dig
  419. assert f64.cast_check(0.5) == Float(0.5, 17)
  420. assert abs(f64.cast_check(3.7) - 3.7) < 3e-17
  421. assert isinstance(f64.cast_check(3), (Float, float))
  422. assert f64.cast_nocheck(oo) == float('inf')
  423. assert f64.cast_nocheck(-oo) == float('-inf')
  424. assert f64.cast_nocheck(float(oo)) == float('inf')
  425. assert f64.cast_nocheck(float(-oo)) == float('-inf')
  426. assert math.isnan(f64.cast_nocheck(nan))
  427. assert f32 != f64
  428. assert f64 == f64.func(*f64.args)
  429. def test_Type__cast_check__floating_point():
  430. raises(ValueError, lambda: f32.cast_check(123.45678949))
  431. raises(ValueError, lambda: f32.cast_check(12.345678949))
  432. raises(ValueError, lambda: f32.cast_check(1.2345678949))
  433. raises(ValueError, lambda: f32.cast_check(.12345678949))
  434. assert abs(123.456789049 - f32.cast_check(123.456789049) - 4.9e-8) < 1e-8
  435. assert abs(0.12345678904 - f32.cast_check(0.12345678904) - 4e-11) < 1e-11
  436. dcm21 = Float('0.123456789012345670499') # 21 decimals
  437. assert abs(dcm21 - f64.cast_check(dcm21) - 4.99e-19) < 1e-19
  438. f80.cast_check(Float('0.12345678901234567890103', precision=88))
  439. raises(ValueError, lambda: f80.cast_check(Float('0.12345678901234567890149', precision=88)))
  440. v10 = 12345.67894
  441. raises(ValueError, lambda: f32.cast_check(v10))
  442. assert abs(Float(str(v10), precision=64+8) - f64.cast_check(v10)) < v10*1e-16
  443. assert abs(f32.cast_check(2147483647) - 2147483650) < 1
  444. def test_Type__cast_check__complex_floating_point():
  445. val9_11 = 123.456789049 + 0.123456789049j
  446. raises(ValueError, lambda: c64.cast_check(.12345678949 + .12345678949j))
  447. assert abs(val9_11 - c64.cast_check(val9_11) - 4.9e-8) < 1e-8
  448. dcm21 = Float('0.123456789012345670499') + 1e-20j # 21 decimals
  449. assert abs(dcm21 - c128.cast_check(dcm21) - 4.99e-19) < 1e-19
  450. v19 = Float('0.1234567890123456749') + 1j*Float('0.1234567890123456749')
  451. raises(ValueError, lambda: c128.cast_check(v19))
  452. def test_While():
  453. xpp = AddAugmentedAssignment(x, 1)
  454. whl1 = While(x < 2, [xpp])
  455. assert whl1.condition.args[0] == x
  456. assert whl1.condition.args[1] == 2
  457. assert whl1.condition == Lt(x, 2, evaluate=False)
  458. assert whl1.body.args == (xpp,)
  459. assert whl1.func(*whl1.args) == whl1
  460. cblk = CodeBlock(AddAugmentedAssignment(x, 1))
  461. whl2 = While(x < 2, cblk)
  462. assert whl1 == whl2
  463. assert whl1 != While(x < 3, [xpp])
  464. def test_Scope():
  465. assign = Assignment(x, y)
  466. incr = AddAugmentedAssignment(x, 1)
  467. scp = Scope([assign, incr])
  468. cblk = CodeBlock(assign, incr)
  469. assert scp.body == cblk
  470. assert scp == Scope(cblk)
  471. assert scp != Scope([incr, assign])
  472. assert scp.func(*scp.args) == scp
  473. def test_Print():
  474. fmt = "%d %.3f"
  475. ps = Print([n, x], fmt)
  476. assert str(ps.format_string) == fmt
  477. assert ps.print_args == Tuple(n, x)
  478. assert ps.args == (Tuple(n, x), QuotedString(fmt), none)
  479. assert ps == Print((n, x), fmt)
  480. assert ps != Print([x, n], fmt)
  481. assert ps.func(*ps.args) == ps
  482. ps2 = Print([n, x])
  483. assert ps2 == Print([n, x])
  484. assert ps2 != ps
  485. assert ps2.format_string == None
  486. def test_FunctionPrototype_and_FunctionDefinition():
  487. vx = Variable(x, type=real)
  488. vn = Variable(n, type=integer)
  489. fp1 = FunctionPrototype(real, 'power', [vx, vn])
  490. assert fp1.return_type == real
  491. assert fp1.name == String('power')
  492. assert fp1.parameters == Tuple(vx, vn)
  493. assert fp1 == FunctionPrototype(real, 'power', [vx, vn])
  494. assert fp1 != FunctionPrototype(real, 'power', [vn, vx])
  495. assert fp1.func(*fp1.args) == fp1
  496. body = [Assignment(x, x**n), Return(x)]
  497. fd1 = FunctionDefinition(real, 'power', [vx, vn], body)
  498. assert fd1.return_type == real
  499. assert str(fd1.name) == 'power'
  500. assert fd1.parameters == Tuple(vx, vn)
  501. assert fd1.body == CodeBlock(*body)
  502. assert fd1 == FunctionDefinition(real, 'power', [vx, vn], body)
  503. assert fd1 != FunctionDefinition(real, 'power', [vx, vn], body[::-1])
  504. assert fd1.func(*fd1.args) == fd1
  505. fp2 = FunctionPrototype.from_FunctionDefinition(fd1)
  506. assert fp2 == fp1
  507. fd2 = FunctionDefinition.from_FunctionPrototype(fp1, body)
  508. assert fd2 == fd1
  509. def test_Return():
  510. rs = Return(x)
  511. assert rs.args == (x,)
  512. assert rs == Return(x)
  513. assert rs != Return(y)
  514. assert rs.func(*rs.args) == rs
  515. def test_FunctionCall():
  516. fc = FunctionCall('power', (x, 3))
  517. assert fc.function_args[0] == x
  518. assert fc.function_args[1] == 3
  519. assert len(fc.function_args) == 2
  520. assert isinstance(fc.function_args[1], Integer)
  521. assert fc == FunctionCall('power', (x, 3))
  522. assert fc != FunctionCall('power', (3, x))
  523. assert fc != FunctionCall('Power', (x, 3))
  524. assert fc.func(*fc.args) == fc
  525. fc2 = FunctionCall('fma', [2, 3, 4])
  526. assert len(fc2.function_args) == 3
  527. assert fc2.function_args[0] == 2
  528. assert fc2.function_args[1] == 3
  529. assert fc2.function_args[2] == 4
  530. assert str(fc2) in ( # not sure if QuotedString is a better default...
  531. 'FunctionCall(fma, function_args=(2, 3, 4))',
  532. 'FunctionCall("fma", function_args=(2, 3, 4))',
  533. )
  534. def test_ast_replace():
  535. x = Variable('x', real)
  536. y = Variable('y', real)
  537. n = Variable('n', integer)
  538. pwer = FunctionDefinition(real, 'pwer', [x, n], [pow(x.symbol, n.symbol)])
  539. pname = pwer.name
  540. pcall = FunctionCall('pwer', [y, 3])
  541. tree1 = CodeBlock(pwer, pcall)
  542. assert str(tree1.args[0].name) == 'pwer'
  543. assert str(tree1.args[1].name) == 'pwer'
  544. for a, b in zip(tree1, [pwer, pcall]):
  545. assert a == b
  546. tree2 = tree1.replace(pname, String('power'))
  547. assert str(tree1.args[0].name) == 'pwer'
  548. assert str(tree1.args[1].name) == 'pwer'
  549. assert str(tree2.args[0].name) == 'power'
  550. assert str(tree2.args[1].name) == 'power'