Add support for complex numbers

Summary:
Here we improve the following:
 * The test case for complex numbers
 * Errpy will now throw a recoverable error when a invalid number is attempted to be parsed (rather than panic crashing)
 * "integer style" and "float style" complex numbers are both handled.
 * The pretty printers are improved to handle complex numbers
 * `ast.rs` and `ast.ml` are both enhanced to support complex numbers

Reviewed By: stroxler

Differential Revision:
D42407324

Privacy Context Container: L1152058

fbshipit-source-id: 5a34ccd8919abf9d01af2c63d46149e93a04efe6

Author: jasontatton

src/ast.ml

@@ -8,6 +8,7 @@
 type num =
   | Int of int
   | Float of float
+  | Complex of float
   | Big_int of string
       (** pyast uses Py.Object.t for Big_int, but this is presumably not what we
        * want (if it's an opaque wrapper for a CPython object). Just have the

src/ast.rs

@@ -17,6 +17,7 @@ use serde::Serialize;
 pub enum Num {
     Int(isize),
     Float(f64),
+    Complex(f64),
     /// pyast uses Py.Object.t for Big_int, but this is presumably not what we
     /// want (if it's an opaque wrapper for a CPython object). Just have the
     /// source text representing the integer digits for now instead.

src/cst_to_ast.rs

@@ -2064,8 +2064,8 @@ impl Parser {
                 KEYWORD_IDENTIFIER => self.name(self.get_text(rule.node)),
                 STRING => self.raw_string(rule.node, rule.node)?,
                 CONCATENATED_STRING => self.concatenated_string(rule.node)?,
-                INTEGER => self.integer(self.get_text(rule.node)),
-                FLOAT => self.float(self.get_text(rule.node)),
+                INTEGER => self.integer(rule.node)?,
+                FLOAT => self.float(rule.node)?,
                 TRUE => self.constant(ConstantDesc::Bool(true)),
                 FALSE => self.constant(ConstantDesc::Bool(false)),
                 NONE => self.none(),
@@ -3426,22 +3426,66 @@ impl Parser {
         }
     }
 
-    fn integer(&mut self, const_value: String) -> ExprDesc {
-        ExprDesc::Constant {
-            value: Some(ConstantDesc::Num(Num::Int(
-                const_value.parse::<isize>().unwrap(),
-            ))),
-            kind: None,
+    fn integer(&mut self, node: &Node) -> ErrorableResult<ExprDesc> {
+        let mut const_value = self.get_text(node);
+
+        let is_complex = const_value.ends_with('j') || const_value.ends_with('J');
+        if is_complex {
+            const_value = const_value[0..const_value.len() - 1].to_string();
         }
-    }
 
-    fn float(&mut self, const_value: String) -> ExprDesc {
-        ExprDesc::Constant {
-            value: Some(ConstantDesc::Num(Num::Float(
-                const_value.parse::<f64>().unwrap(),
-            ))),
+        let integer_value = match const_value.parse::<isize>() {
+            Ok(value) => value,
+            Err(error_msg) => {
+                return Err(self.record_recoverable_error(
+                    RecoverableError::UnexpectedExpression(format!(
+                        "cannot parse integer: {:?} as {:?}",
+                        const_value, error_msg
+                    )),
+                    node,
+                ));
+            }
+        };
+
+        Ok(ExprDesc::Constant {
+            value: Some(ConstantDesc::Num(if is_complex {
+                Num::Complex(integer_value as f64)
+            } else {
+                Num::Int(integer_value)
+            })),
             kind: None,
+        })
+    }
+
+    fn float(&mut self, node: &Node) -> ErrorableResult<ExprDesc> {
+        let mut const_value = self.get_text(node);
+
+        let is_complex = const_value.ends_with('j') || const_value.ends_with('J');
+        if is_complex {
+            const_value = const_value[0..const_value.len() - 1].to_string();
         }
+
+        let float_value = match const_value.parse::<f64>() {
+            Ok(value) => value,
+            Err(error_msg) => {
+                return Err(self.record_recoverable_error(
+                    RecoverableError::UnexpectedExpression(format!(
+                        "cannot parse float: {:?} as {:?}",
+                        const_value, error_msg
+                    )),
+                    node,
+                ));
+            }
+        };
+
+        Ok(ExprDesc::Constant {
+            value: Some(ConstantDesc::Num(if is_complex {
+                Num::Complex(float_value)
+            } else {
+                Num::Float(float_value)
+            })),
+            kind: None,
+        })
     }
 
     fn name(&mut self, identifier: String) -> ExprDesc {

src/printers/ast_pretty_print.rs

@@ -1260,6 +1260,10 @@ impl ConstantDesc {
             ConstantDesc::Num(Num::Float(vala)) => {
                 pprint_output.push_str(&vala.to_string());
             }
+            ConstantDesc::Num(Num::Complex(vala)) => {
+                pprint_output.push_str(&vala.to_string());
+                pprint_output.push_str("j");
+            }
             rest => {
                 let res = match rest {
                     ConstantDesc::Str(vala) => vala,

src/printers/ast_print.rs

@@ -749,6 +749,9 @@ impl fmt::Display for ConstantDesc {
             ConstantDesc::Num(Num::Float(vala)) => {
                 write!(f, "{}", vala)
             }
+            ConstantDesc::Num(Num::Complex(vala)) => {
+                write!(f, "{}j", vala)
+            }
             rest => {
                 let res = match rest {
                     ConstantDesc::Str(vala) => vala,

tests/test_resources/unit_tests/ast_literals.pytest

@@ -234,7 +234,14 @@ f"{{{{70 + 4}}}}"
 timestamp = 9223372036854775810
 
 ## complex numbers
-complex = 0.0302988j
+complex_half = 0.0302988j
+
+c = 1 + 2.j
+c = 1 + 2j
+c = 1 + 2.J
+c = 1 + 2J
+real = c.real
+imag = c.imag
 
 ## very long floats
 timestamp = 0.33333333333333334444444444444444356456456456456456457457456456456456456456

tests/test_resources/unit_tests/ast_literals.pytest.expect_fails

@@ -1,2 +1 @@
 long integers
-complex numbers