bpo-32894: Support unparsing of infinity numbers in ast_unparser.c (GH-17426)

3 files changed, 68 insertions, 5 deletions

diff --git a/Lib/test/test_future.py b/Lib/test/test_future.py
index 9b88e3f..ebeb833 100644
--- a/Lib/test/test_future.py
+++ b/Lib/test/test_future.py
@@ -5,6 +5,7 @@ from test import support
 from textwrap import dedent
 import os
 import re
+import sys
 
 rx = re.compile(r'\((\S+).py, line (\d+)')
 
@@ -308,5 +309,18 @@ class AnnotationsFutureTestCase(unittest.TestCase):
         self.assertAnnotationEqual("f'{x=!a}'", expected="f'x={x!a}'")
         self.assertAnnotationEqual("f'{x=!s:*^20}'", expected="f'x={x!s:*^20}'")
 
+    def test_infinity_numbers(self):
+        inf = "1e" + repr(sys.float_info.max_10_exp + 1)
+        infj = f"{inf}j"
+        self.assertAnnotationEqual("1e1000", expected=inf)
+        self.assertAnnotationEqual("1e1000j", expected=infj)
+        self.assertAnnotationEqual("-1e1000", expected=f"-{inf}")
+        self.assertAnnotationEqual("3+1e1000j", expected=f"3 + {infj}")
+        self.assertAnnotationEqual("(1e1000, 1e1000j)", expected=f"({inf}, {infj})")
+        self.assertAnnotationEqual("'inf'")
+        self.assertAnnotationEqual("('inf', 1e1000, 'infxxx', 1e1000j)", expected=f"('inf', {inf}, 'infxxx', {infj})")
+        self.assertAnnotationEqual("(1e1000, (1e1000j,))", expected=f"({inf}, ({infj},))")
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2019-12-01-21-36-49.bpo-32894.5g_UQr.rst b/Misc/NEWS.d/next/Core and Builtins/2019-12-01-21-36-49.bpo-32894.5g_UQr.rst
new file mode 100644
index 0000000..68f4e67
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2019-12-01-21-36-49.bpo-32894.5g_UQr.rst
@@ -0,0 +1 @@
+Support unparsing of infinity numbers in postponed annotations. Patch by Batuhan Taşkaya.
diff --git a/Python/ast_unparse.c b/Python/ast_unparse.c
index 5ecd1b0..7cf199b 100644
--- a/Python/ast_unparse.c
+++ b/Python/ast_unparse.c
@@ -1,3 +1,4 @@
+#include <float.h>   /* DBL_MAX_10_EXP */
 #include <stdbool.h>
 #include "Python.h"
 #include "Python-ast.h"
@@ -6,6 +7,8 @@ static PyObject *_str_open_br;
 static PyObject *_str_dbl_open_br;
 static PyObject *_str_close_br;
 static PyObject *_str_dbl_close_br;
+static PyObject *_str_inf;
+static PyObject *_str_replace_inf;
 
 /* Forward declarations for recursion via helper functions. */
 static PyObject *
@@ -61,13 +64,28 @@ append_charp(_PyUnicodeWriter *writer, const char *charp)
 static int
 append_repr(_PyUnicodeWriter *writer, PyObject *obj)
 {
-    int ret;
-    PyObject *repr;
-    repr = PyObject_Repr(obj);
+    PyObject *repr = PyObject_Repr(obj);
+
     if (!repr) {
         return -1;
     }
-    ret = _PyUnicodeWriter_WriteStr(writer, repr);
+
+    if ((PyFloat_CheckExact(obj) && Py_IS_INFINITY(PyFloat_AS_DOUBLE(obj))) ||
+       PyComplex_CheckExact(obj))
+    {
+        PyObject *new_repr = PyUnicode_Replace(
+            repr,
+            _str_inf,
+            _str_replace_inf,
+            -1
+        );
+        Py_DECREF(repr);
+        if (!new_repr) {
+            return -1;
+        }
+        repr = new_repr;
+    }
+    int ret = _PyUnicodeWriter_WriteStr(writer, repr);
     Py_DECREF(repr);
     return ret;
 }
@@ -698,6 +716,28 @@ append_formattedvalue(_PyUnicodeWriter *writer, expr_ty e)
 }
 
 static int
+append_ast_constant(_PyUnicodeWriter *writer, PyObject *constant)
+{
+    if (PyTuple_CheckExact(constant)) {
+        Py_ssize_t i, elem_count;
+
+        elem_count = PyTuple_GET_SIZE(constant);
+        APPEND_STR("(");
+        for (i = 0; i < elem_count; i++) {
+            APPEND_STR_IF(i > 0, ", ");
+            if (append_ast_constant(writer, PyTuple_GET_ITEM(constant, i)) < 0) {
+                return -1;
+            }
+        }
+
+        APPEND_STR_IF(elem_count == 1, ",");
+        APPEND_STR(")");
+        return 0;
+    }
+    return append_repr(writer, constant);
+}
+
+static int
 append_ast_attribute(_PyUnicodeWriter *writer, expr_ty e)
 {
     const char *period;
@@ -835,7 +875,7 @@ append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level)
         if (e->v.Constant.value == Py_Ellipsis) {
             APPEND_STR_FINISH("...");
         }
-        return append_repr(writer, e->v.Constant.value);
+        return append_ast_constant(writer, e->v.Constant.value);
     case JoinedStr_kind:
         return append_joinedstr(writer, e, false);
     case FormattedValue_kind:
@@ -883,6 +923,14 @@ maybe_init_static_strings(void)
         !(_str_dbl_close_br = PyUnicode_InternFromString("}}"))) {
         return -1;
     }
+    if (!_str_inf &&
+        !(_str_inf = PyUnicode_FromString("inf"))) {
+        return -1;
+    }
+    if (!_str_replace_inf &&
+        !(_str_replace_inf = PyUnicode_FromFormat("1e%d", 1 + DBL_MAX_10_EXP))) {
+        return -1;
+    }
     return 0;
 }