Coverage for /usr/lib/python3/dist-packages/sympy/interactive/traversal.py: 6%

1from sympy.core.basic import Basic 

2from sympy.printing import pprint 

3 

4import random 

5 

6def interactive_traversal(expr): 

7 """Traverse a tree asking a user which branch to choose. """ 

8 

9 RED, BRED = '\033[0;31m', '\033[1;31m' 

10 GREEN, BGREEN = '\033[0;32m', '\033[1;32m' 

11 YELLOW, BYELLOW = '\033[0;33m', '\033[1;33m' # noqa 

12 BLUE, BBLUE = '\033[0;34m', '\033[1;34m' # noqa 

13 MAGENTA, BMAGENTA = '\033[0;35m', '\033[1;35m'# noqa 

14 CYAN, BCYAN = '\033[0;36m', '\033[1;36m' # noqa 

15 END = '\033[0m' 

16 

17 def cprint(*args): 

18 print("".join(map(str, args)) + END) 

19 

20 def _interactive_traversal(expr, stage): 

21 if stage > 0: 

22 print() 

23 

24 cprint("Current expression (stage ", BYELLOW, stage, END, "):") 

25 print(BCYAN) 

26 pprint(expr) 

27 print(END) 

28 

29 if isinstance(expr, Basic): 

30 if expr.is_Add: 

31 args = expr.as_ordered_terms() 

32 elif expr.is_Mul: 

33 args = expr.as_ordered_factors() 

34 else: 

35 args = expr.args 

36 elif hasattr(expr, "__iter__"): 

37 args = list(expr) 

38 else: 

39 return expr 

40 

41 n_args = len(args) 

42 

43 if not n_args: 

44 return expr 

45 

46 for i, arg in enumerate(args): 

47 cprint(GREEN, "[", BGREEN, i, GREEN, "] ", BLUE, type(arg), END) 

48 pprint(arg) 

49 print() 

50 

51 if n_args == 1: 

52 choices = '0' 

53 else: 

54 choices = '0-%d' % (n_args - 1) 

55 

56 try: 

57 choice = input("Your choice [%s,f,l,r,d,?]: " % choices) 

58 except EOFError: 

59 result = expr 

60 print() 

61 else: 

62 if choice == '?': 

63 cprint(RED, "%s - select subexpression with the given index" % 

64 choices) 

65 cprint(RED, "f - select the first subexpression") 

66 cprint(RED, "l - select the last subexpression") 

67 cprint(RED, "r - select a random subexpression") 

68 cprint(RED, "d - done\n") 

69 

70 result = _interactive_traversal(expr, stage) 

71 elif choice in ('d', ''): 

72 result = expr 

73 elif choice == 'f': 

74 result = _interactive_traversal(args[0], stage + 1) 

75 elif choice == 'l': 

76 result = _interactive_traversal(args[-1], stage + 1) 

77 elif choice == 'r': 

78 result = _interactive_traversal(random.choice(args), stage + 1) 

79 else: 

80 try: 

81 choice = int(choice) 

82 except ValueError: 

83 cprint(BRED, 

84 "Choice must be a number in %s range\n" % choices) 

85 result = _interactive_traversal(expr, stage) 

86 else: 

87 if choice < 0 or choice >= n_args: 

88 cprint(BRED, "Choice must be in %s range\n" % choices) 

89 result = _interactive_traversal(expr, stage) 

90 else: 

91 result = _interactive_traversal(args[choice], stage + 1) 

92 

93 return result 

94 

95 return _interactive_traversal(expr, 0)