get_type(...) can now guess the return type of binary operators

  - including boolean operators and comparison operators

1 files changed, 13 insertions, 0 deletions

diff --git a/TurtleArt/tatype.py b/TurtleArt/tatype.py
index 055fc2f..cf7bb03 100644
--- a/TurtleArt/tatype.py
+++ b/TurtleArt/tatype.py
@@ -131,6 +131,19 @@ def get_type(x):
     # unary operands never change the type of their argument
     elif isinstance(x, ast.UnaryOp):
         return get_type(x.operand)
+    # boolean and comparison operators always return a boolean
+    if isinstance(x, (ast.BoolOp, ast.Compare)):
+        return (TYPE_BOOL, True)
+    # other binary operators
+    elif isinstance(x, ast.BinOp):
+        type_left = get_type(x.left)[0]
+        if type_left == TYPE_STRING:
+            return (TYPE_STRING, True)
+        type_right = get_type(x.right)[0]
+        if type_left == type_right == TYPE_INT:
+            return (TYPE_INT, True)
+        else:
+            return (TYPE_FLOAT, True)
 
     return (TYPE_OBJECT, isinstance(x, ast.AST))