Spaces:
Sleeping
Sleeping
File size: 3,189 Bytes
6a86ad5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
from sympy.core.basic import Basic
from sympy.printing import pprint
import random
def interactive_traversal(expr):
"""Traverse a tree asking a user which branch to choose. """
RED, BRED = '\033[0;31m', '\033[1;31m'
GREEN, BGREEN = '\033[0;32m', '\033[1;32m'
YELLOW, BYELLOW = '\033[0;33m', '\033[1;33m' # noqa
BLUE, BBLUE = '\033[0;34m', '\033[1;34m' # noqa
MAGENTA, BMAGENTA = '\033[0;35m', '\033[1;35m'# noqa
CYAN, BCYAN = '\033[0;36m', '\033[1;36m' # noqa
END = '\033[0m'
def cprint(*args):
print("".join(map(str, args)) + END)
def _interactive_traversal(expr, stage):
if stage > 0:
print()
cprint("Current expression (stage ", BYELLOW, stage, END, "):")
print(BCYAN)
pprint(expr)
print(END)
if isinstance(expr, Basic):
if expr.is_Add:
args = expr.as_ordered_terms()
elif expr.is_Mul:
args = expr.as_ordered_factors()
else:
args = expr.args
elif hasattr(expr, "__iter__"):
args = list(expr)
else:
return expr
n_args = len(args)
if not n_args:
return expr
for i, arg in enumerate(args):
cprint(GREEN, "[", BGREEN, i, GREEN, "] ", BLUE, type(arg), END)
pprint(arg)
print()
if n_args == 1:
choices = '0'
else:
choices = '0-%d' % (n_args - 1)
try:
choice = input("Your choice [%s,f,l,r,d,?]: " % choices)
except EOFError:
result = expr
print()
else:
if choice == '?':
cprint(RED, "%s - select subexpression with the given index" %
choices)
cprint(RED, "f - select the first subexpression")
cprint(RED, "l - select the last subexpression")
cprint(RED, "r - select a random subexpression")
cprint(RED, "d - done\n")
result = _interactive_traversal(expr, stage)
elif choice in ('d', ''):
result = expr
elif choice == 'f':
result = _interactive_traversal(args[0], stage + 1)
elif choice == 'l':
result = _interactive_traversal(args[-1], stage + 1)
elif choice == 'r':
result = _interactive_traversal(random.choice(args), stage + 1)
else:
try:
choice = int(choice)
except ValueError:
cprint(BRED,
"Choice must be a number in %s range\n" % choices)
result = _interactive_traversal(expr, stage)
else:
if choice < 0 or choice >= n_args:
cprint(BRED, "Choice must be in %s range\n" % choices)
result = _interactive_traversal(expr, stage)
else:
result = _interactive_traversal(args[choice], stage + 1)
return result
return _interactive_traversal(expr, 0)
|