Support plus on numbers in PyEvaluator (PY-19856)

python/src/com/jetbrains/python/psi/impl/PyEvaluator.java

@@ -106,17 +106,7 @@ public class PyEvaluator {
     if (expression.isIntegerLiteral()) {
       final BigInteger value = expression.getBigIntegerValue();
       if (value != null) {
-        final int intValue = value.intValue();
-        if (BigInteger.valueOf(intValue).equals(value)) {
-          return intValue;
-        }
-
-        final long longValue = value.longValue();
-        if (BigInteger.valueOf(longValue).equals(value)) {
-          return longValue;
-        }
-
-        return value;
+        return fromBigInteger(value);
       }
     }
 
@@ -182,6 +172,11 @@ public class PyEvaluator {
     else if (lhs instanceof List && rhs instanceof List) {
       return ContainerUtil.concat((List)lhs, (List)rhs);
     }
+    else if (lhs instanceof Number && rhs instanceof Number) {
+      final BigInteger first = toBigInteger((Number)lhs);
+      final BigInteger second = toBigInteger((Number)rhs);
+      return fromBigInteger(first.add(second));
+    }
     return null;
   }
 
@@ -246,6 +241,34 @@ public class PyEvaluator {
     result.put(myEvaluateKeys ? evaluate(key) : key, myEvaluateCollectionItems ? evaluate(value) : value);
   }
 
+  @NotNull
+  private static BigInteger toBigInteger(@NotNull Number value) {
+    // don't forget to update fromBigInteger() after changing this method
+
+    return value instanceof Integer
+           ? BigInteger.valueOf((Integer)value)
+           : value instanceof Long
+             ? BigInteger.valueOf((Long)value)
+             : (BigInteger)value;
+  }
+
+  @NotNull
+  private static Number fromBigInteger(@NotNull BigInteger value) {
+    // don't forget to update toBigInteger() after changing this method
+
+    final int intValue = value.intValue();
+    if (BigInteger.valueOf(intValue).equals(value)) {
+      return intValue;
+    }
+
+    final long longValue = value.longValue();
+    if (BigInteger.valueOf(longValue).equals(value)) {
+      return longValue;
+    }
+
+    return value;
+  }
+
   /**
    * Shortcut that evaluates expression with default params and casts it to particular type (if possible)
    *

python/testSrc/com/jetbrains/python/PyEvaluatorTest.java

@@ -224,6 +224,16 @@ public class PyEvaluatorTest extends PyTestCase {
     assertNull(evaluator.evaluate(parseText("a = True\nb = a\nc = b\nexpr = c")));
   }
 
+  public void testNumbersAddition() {
+    assertEquals(Integer.valueOf(3), byExpression("1 + 2", Integer.class));
+
+    assertEquals(Long.valueOf(Long.valueOf(Integer.MAX_VALUE) + 1), byExpression("1 + " + Integer.MAX_VALUE, Long.class));
+    assertEquals(Long.valueOf(Long.valueOf(Integer.MAX_VALUE) + 1), byExpression(Integer.MAX_VALUE + "+ 1", Long.class));
+
+    assertEquals(BigInteger.valueOf(Long.MAX_VALUE).add(BigInteger.ONE), byExpression("1 + " + Long.MAX_VALUE, BigInteger.class));
+    assertEquals(BigInteger.valueOf(Long.MAX_VALUE).add(BigInteger.ONE), byExpression(Long.MAX_VALUE + "+ 1", BigInteger.class));
+  }
+
   @NotNull
   private <T> T byExpression(@NotNull String expression, @NotNull Class<T> cls) {
     final Object value = new PyEvaluator().evaluate(parseExpression(expression));