| 386 | assert fromstring("x .ge. y", language=Language.Fortran) == as_ge(x, y) |
| 387 | |
| 388 | def test_traverse(self): |
| 389 | x = as_symbol("x") |
| 390 | y = as_symbol("y") |
| 391 | z = as_symbol("z") |
| 392 | f = as_symbol("f") |
| 393 | |
| 394 | # Use traverse to substitute a symbol |
| 395 | def replace_visit(s, r=z): |
| 396 | if s == x: |
| 397 | return r |
| 398 | |
| 399 | assert x.traverse(replace_visit) == z |
| 400 | assert y.traverse(replace_visit) == y |
| 401 | assert z.traverse(replace_visit) == z |
| 402 | assert (f(y)).traverse(replace_visit) == f(y) |
| 403 | assert (f(x)).traverse(replace_visit) == f(z) |
| 404 | assert (f[y]).traverse(replace_visit) == f[y] |
| 405 | assert (f[z]).traverse(replace_visit) == f[z] |
| 406 | assert (x + y + z).traverse(replace_visit) == (2 * z + y) |
| 407 | assert (x + |
| 408 | f(y, x - z)).traverse(replace_visit) == (z + |
| 409 | f(y, as_number(0))) |
| 410 | assert as_eq(x, y).traverse(replace_visit) == as_eq(z, y) |
| 411 | |
| 412 | # Use traverse to collect symbols, method 1 |
| 413 | function_symbols = set() |
| 414 | symbols = set() |
| 415 | |
| 416 | def collect_symbols(s): |
| 417 | if s.op is Op.APPLY: |
| 418 | oper = s.data[0] |
| 419 | function_symbols.add(oper) |
| 420 | if oper in symbols: |
| 421 | symbols.remove(oper) |
| 422 | elif s.op is Op.SYMBOL and s not in function_symbols: |
| 423 | symbols.add(s) |
| 424 | |
| 425 | (x + f(y, x - z)).traverse(collect_symbols) |
| 426 | assert function_symbols == {f} |
| 427 | assert symbols == {x, y, z} |
| 428 | |
| 429 | # Use traverse to collect symbols, method 2 |
| 430 | def collect_symbols2(expr, symbols): |
| 431 | if expr.op is Op.SYMBOL: |
| 432 | symbols.add(expr) |
| 433 | |
| 434 | symbols = set() |
| 435 | (x + f(y, x - z)).traverse(collect_symbols2, symbols) |
| 436 | assert symbols == {x, y, z, f} |
| 437 | |
| 438 | # Use traverse to partially collect symbols |
| 439 | def collect_symbols3(expr, symbols): |
| 440 | if expr.op is Op.APPLY: |
| 441 | # skip traversing function calls |
| 442 | return expr |
| 443 | if expr.op is Op.SYMBOL: |
| 444 | symbols.add(expr) |
| 445 | |