Issue #1588: Add complex.__format__.

6 files changed, 426 insertions, 53 deletions

diff --git a/Include/complexobject.h b/Include/complexobject.h
index 84b6d8b7e1..8290a6dc4f 100644
--- a/Include/complexobject.h
+++ b/Include/complexobject.h
@@ -54,6 +54,12 @@ PyAPI_FUNC(double) PyComplex_RealAsDouble(PyObject *op);
 PyAPI_FUNC(double) PyComplex_ImagAsDouble(PyObject *op);
 PyAPI_FUNC(Py_complex) PyComplex_AsCComplex(PyObject *op);
 
+/* Format the object based on the format_spec, as defined in PEP 3101
+   (Advanced String Formatting). */
+PyAPI_FUNC(PyObject *) _PyComplex_FormatAdvanced(PyObject *obj,
+                                                 Py_UNICODE *format_spec,
+                                                 Py_ssize_t format_spec_len);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py
index 87e25841fc..ac19353642 100644
--- a/Lib/test/test_complex.py
+++ b/Lib/test/test_complex.py
@@ -436,7 +436,66 @@ class ComplexTest(unittest.TestCase):
                 self.assertFloatsAreIdentical(0.0 + z.imag,
                                               0.0 + roundtrip.imag)
 
-
+    def test_format(self):
+        # empty format string is same as str()
+        self.assertEqual(format(1+3j, ''), str(1+3j))
+        self.assertEqual(format(1.5+3.5j, ''), str(1.5+3.5j))
+        self.assertEqual(format(3j, ''), str(3j))
+        self.assertEqual(format(3.2j, ''), str(3.2j))
+        self.assertEqual(format(3+0j, ''), str(3+0j))
+        self.assertEqual(format(3.2+0j, ''), str(3.2+0j))
+
+        self.assertEqual(format(1+3j, 'g'), '1+3j')
+        self.assertEqual(format(3j, 'g'), '0+3j')
+        self.assertEqual(format(1.5+3.5j, 'g'), '1.5+3.5j')
+
+        self.assertEqual(format(1.5+3.5j, '+g'), '+1.5+3.5j')
+        self.assertEqual(format(1.5-3.5j, '+g'), '+1.5-3.5j')
+        self.assertEqual(format(1.5-3.5j, '-g'), '1.5-3.5j')
+        self.assertEqual(format(1.5+3.5j, ' g'), ' 1.5+3.5j')
+        self.assertEqual(format(1.5-3.5j, ' g'), ' 1.5-3.5j')
+        self.assertEqual(format(-1.5+3.5j, ' g'), '-1.5+3.5j')
+        self.assertEqual(format(-1.5-3.5j, ' g'), '-1.5-3.5j')
+
+        self.assertEqual(format(-1.5-3.5e-20j, 'g'), '-1.5-3.5e-20j')
+        self.assertEqual(format(-1.5-3.5j, 'f'), '-1.500000-3.500000j')
+        self.assertEqual(format(-1.5-3.5j, 'F'), '-1.500000-3.500000j')
+        self.assertEqual(format(-1.5-3.5j, 'e'), '-1.500000e+00-3.500000e+00j')
+        self.assertEqual(format(-1.5-3.5j, '.2e'), '-1.50e+00-3.50e+00j')
+        self.assertEqual(format(-1.5-3.5j, '.2E'), '-1.50E+00-3.50E+00j')
+        self.assertEqual(format(-1.5e10-3.5e5j, '.2G'), '-1.5E+10-3.5E+05j')
+
+        self.assertEqual(format(1.5+3j, '<20g'),  '1.5+3j              ')
+        self.assertEqual(format(1.5+3j, '*<20g'), '1.5+3j**************')
+        self.assertEqual(format(1.5+3j, '>20g'),  '              1.5+3j')
+        self.assertEqual(format(1.5+3j, '^20g'),  '       1.5+3j       ')
+        self.assertEqual(format(1.5+3j, '<20'),   '(1.5+3j)            ')
+        self.assertEqual(format(1.5+3j, '>20'),   '            (1.5+3j)')
+        self.assertEqual(format(1.5+3j, '^20'),   '      (1.5+3j)      ')
+        self.assertEqual(format(1.123-3.123j, '^20.2'), '     (1.1-3.1j)     ')
+
+        self.assertEqual(format(1.5+3j, '<20.2f'), '1.50+3.00j          ')
+        self.assertEqual(format(1.5e20+3j, '<20.2f'), '150000000000000000000.00+3.00j')
+        self.assertEqual(format(1.5e20+3j, '>40.2f'), '          150000000000000000000.00+3.00j')
+        self.assertEqual(format(1.5e20+3j, '^40,.2f'), '  150,000,000,000,000,000,000.00+3.00j  ')
+        self.assertEqual(format(1.5e21+3j, '^40,.2f'), ' 1,500,000,000,000,000,000,000.00+3.00j ')
+        self.assertEqual(format(1.5e21+3000j, ',.2f'), '1,500,000,000,000,000,000,000.00+3,000.00j')
+
+        # alternate is invalid
+        self.assertRaises(ValueError, (1.5+0.5j).__format__, '#f')
+
+        # zero padding is invalid
+        self.assertRaises(ValueError, (1.5+0.5j).__format__, '010f')
+
+        # '=' alignment is invalid
+        self.assertRaises(ValueError, (1.5+3j).__format__, '=20')
+
+        # integer presentation types are an error
+        for t in 'bcdoxX':
+            self.assertRaises(ValueError, (1.5+0.5j).__format__, t)
+
+        # make sure everything works in ''.format()
+        self.assertEqual('*{0:.3f}*'.format(3.14159+2.71828j), '*3.142+2.718j*')
 
 def test_main():
     support.run_unittest(ComplexTest)
diff --git a/Misc/NEWS b/Misc/NEWS
index 01313f8b61..a14682a47c 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -12,6 +12,9 @@ What's New in Python 3.1 beta 1?
 Core and Builtins
 -----------------
 
+- Issue #1588: Add complex.__format__. For example, 
+  format(complex(1, 2./3), '.5') now produces a sensible result.
+
 - Issue #5864: Fix empty format code formatting for floats so that it
   never gives more than the requested number of significant digits.
 
diff --git a/Objects/complexobject.c b/Objects/complexobject.c
index acd885b459..691809f134 100644
--- a/Objects/complexobject.c
+++ b/Objects/complexobject.c
@@ -681,6 +681,23 @@ complex_getnewargs(PyComplexObject *v)
 	return Py_BuildValue("(dd)", c.real, c.imag);
 }
 
+PyDoc_STRVAR(complex__format__doc,
+"complex.__format__() -> str\n"
+"\n"
+"Converts to a string according to format_spec.");
+
+static PyObject *
+complex__format__(PyObject* self, PyObject* args)
+{
+    PyObject *format_spec;
+
+    if (!PyArg_ParseTuple(args, "U:__format__", &format_spec))
+        return NULL;
+    return _PyComplex_FormatAdvanced(self,
+                                     PyUnicode_AS_UNICODE(format_spec),
+                                     PyUnicode_GET_SIZE(format_spec));
+}
+
 #if 0
 static PyObject *
 complex_is_finite(PyObject *self)
@@ -705,6 +722,8 @@ static PyMethodDef complex_methods[] = {
 	 complex_is_finite_doc},
 #endif
 	{"__getnewargs__",	(PyCFunction)complex_getnewargs,	METH_NOARGS},
+	{"__format__",          (PyCFunction)complex__format__,
+                                           METH_VARARGS, complex__format__doc},
 	{NULL,		NULL}		/* sentinel */
 };
 
diff --git a/Objects/stringlib/formatter.h b/Objects/stringlib/formatter.h
index 9cbd2cc931..c5cda4cebf 100644
--- a/Objects/stringlib/formatter.h
+++ b/Objects/stringlib/formatter.h
@@ -11,6 +11,7 @@
    FORMAT_STRING
    FORMAT_LONG
    FORMAT_FLOAT
+   FORMAT_COMPLEX
    to be whatever you want the public names of these functions to
    be.  These are the only non-static functions defined here.
 */
@@ -261,7 +262,54 @@ parse_internal_render_format_spec(STRINGLIB_CHAR *format_spec,
     return 1;
 }
 
-#if defined FORMAT_FLOAT || defined FORMAT_LONG
+/* Calculate the padding needed. */
+static void
+calc_padding(Py_ssize_t nchars, Py_ssize_t width, STRINGLIB_CHAR align,
+             Py_ssize_t *n_lpadding, Py_ssize_t *n_rpadding,
+             Py_ssize_t *n_total)
+{
+    if (width >= 0) {
+        if (nchars > width)
+            *n_total = nchars;
+        else
+            *n_total = width;
+    }
+    else {
+        /* not specified, use all of the chars and no more */
+        *n_total = nchars;
+    }
+
+    /* figure out how much leading space we need, based on the
+       aligning */
+    if (align == '>')
+        *n_lpadding = *n_total - nchars;
+    else if (align == '^')
+        *n_lpadding = (*n_total - nchars) / 2;
+    else
+        *n_lpadding = 0;
+
+    *n_rpadding = *n_total - nchars - *n_lpadding;
+}
+
+/* Do the padding, and return a pointer to where the caller-supplied
+   content goes. */
+static STRINGLIB_CHAR *
+fill_padding(STRINGLIB_CHAR *p, Py_ssize_t nchars, STRINGLIB_CHAR fill_char,
+             Py_ssize_t n_lpadding, Py_ssize_t n_rpadding)
+{
+    /* Pad on left. */
+    if (n_lpadding)
+        STRINGLIB_FILL(p, fill_char, n_lpadding);
+
+    /* Pad on right. */
+    if (n_rpadding)
+        STRINGLIB_FILL(p + nchars + n_lpadding, fill_char, n_rpadding);
+
+    /* Pointer to the user content. */
+    return p + n_lpadding;
+}
+
+#if defined FORMAT_FLOAT || defined FORMAT_LONG || defined FORMAT_COMPLEX
 /************************************************************************/
 /*********** common routines for numeric formatting *********************/
 /************************************************************************/
@@ -304,6 +352,7 @@ typedef struct {
                                the n_grouped_digits width. */
 } NumberFieldWidths;
 
+
 /* Given a number of the form:
    digits[remainder]
    where ptr points to the start and end points to the end, find where
@@ -564,7 +613,7 @@ get_locale_info(int type, LocaleInfo *locale_info)
     }
 }
 
-#endif /* FORMAT_FLOAT || FORMAT_LONG */
+#endif /* FORMAT_FLOAT || FORMAT_LONG || FORMAT_COMPLEX */
 
 /************************************************************************/
 /*********** string formatting ******************************************/
@@ -573,10 +622,10 @@ get_locale_info(int type, LocaleInfo *locale_info)
 static PyObject *
 format_string_internal(PyObject *value, const InternalFormatSpec *format)
 {
-    Py_ssize_t width; /* total field width */
     Py_ssize_t lpad;
-    STRINGLIB_CHAR *dst;
-    STRINGLIB_CHAR *src = STRINGLIB_STR(value);
+    Py_ssize_t rpad;
+    Py_ssize_t total;
+    STRINGLIB_CHAR *p;
     Py_ssize_t len = STRINGLIB_LEN(value);
     PyObject *result = NULL;
 
@@ -609,56 +658,20 @@ format_string_internal(PyObject *value, const InternalFormatSpec *format)
         len = format->precision;
     }
 
-    if (format->width >= 0) {
-        width = format->width;
-
-        /* but use at least len characters */
-        if (len > width) {
-            width = len;
-        }
-    }
-    else {
-        /* not specified, use all of the chars and no more */
-        width = len;
-    }
+    calc_padding(len, format->width, format->align, &lpad, &rpad, &total);
 
     /* allocate the resulting string */
-    result = STRINGLIB_NEW(NULL, width);
+    result = STRINGLIB_NEW(NULL, total);
     if (result == NULL)
         goto done;
 
-    /* now write into that space */
-    dst = STRINGLIB_STR(result);
-
-    /* figure out how much leading space we need, based on the
-       aligning */
-    if (format->align == '>')
-        lpad = width - len;
-    else if (format->align == '^')
-        lpad = (width - len) / 2;
-    else
-        lpad = 0;
-
-    /* if right aligning, increment the destination allow space on the
-       left */
-    memcpy(dst + lpad, src, len * sizeof(STRINGLIB_CHAR));
-
-    /* do any padding */
-    if (width > len) {
-        STRINGLIB_CHAR fill_char = format->fill_char;
-        if (fill_char == '\0') {
-            /* use the default, if not specified */
-            fill_char = ' ';
-        }
-
-        /* pad on left */
-        if (lpad)
-            STRINGLIB_FILL(dst, fill_char, lpad);
+    /* Write into that space. First the padding. */
+    p = fill_padding(STRINGLIB_STR(result), len,
+                     format->fill_char=='\0'?' ':format->fill_char,
+                     lpad, rpad);
 
-        /* pad on right */
-        if (width - len - lpad)
-            STRINGLIB_FILL(dst + len + lpad, fill_char, width - len - lpad);
-    }
+    /* Then the source string. */
+    memcpy(p, STRINGLIB_STR(value), len * sizeof(STRINGLIB_CHAR));
 
 done:
     return result;
@@ -998,6 +1011,231 @@ done:
 #endif /* FORMAT_FLOAT */
 
 /************************************************************************/
+/*********** complex formatting *****************************************/
+/************************************************************************/
+
+#ifdef FORMAT_COMPLEX
+
+static PyObject *
+format_complex_internal(PyObject *value,
+                        const InternalFormatSpec *format)
+{
+    double re;
+    double im;
+    char *re_buf = NULL;       /* buffer returned from PyOS_double_to_string */
+    char *im_buf = NULL;       /* buffer returned from PyOS_double_to_string */
+
+    InternalFormatSpec tmp_format = *format;
+    Py_ssize_t n_re_digits;
+    Py_ssize_t n_im_digits;
+    Py_ssize_t n_re_remainder;
+    Py_ssize_t n_im_remainder;
+    Py_ssize_t n_re_total;
+    Py_ssize_t n_im_total;
+    int re_has_decimal;
+    int im_has_decimal;
+    Py_ssize_t precision = format->precision;
+    STRINGLIB_CHAR type = format->type;
+    STRINGLIB_CHAR *p_re;
+    STRINGLIB_CHAR *p_im;
+    NumberFieldWidths re_spec;
+    NumberFieldWidths im_spec;
+    int flags = 0;
+    PyObject *result = NULL;
+    STRINGLIB_CHAR *p;
+    STRINGLIB_CHAR re_sign_char = '\0';
+    STRINGLIB_CHAR im_sign_char = '\0';
+    int re_float_type; /* Used to see if we have a nan, inf, or regular float. */
+    int im_float_type;
+    int add_parens = 0;
+    int skip_re = 0;
+    Py_ssize_t lpad;
+    Py_ssize_t rpad;
+    Py_ssize_t total;
+
+#if STRINGLIB_IS_UNICODE
+    Py_UNICODE *re_unicode_tmp = NULL;
+    Py_UNICODE *im_unicode_tmp = NULL;
+#endif
+
+    /* Locale settings, either from the actual locale or
+       from a hard-code pseudo-locale */
+    LocaleInfo locale;
+
+    /* Alternate is not allowed on complex. */
+    if (format->alternate) {
+        PyErr_SetString(PyExc_ValueError,
+                        "Alternate form (#) not allowed in complex format "
+                        "specifier");
+        goto done;
+    }
+
+    /* Neither is zero pading. */
+    if (format->fill_char == '0') {
+        PyErr_SetString(PyExc_ValueError,
+                        "Zero padding is not allowed in complex format "
+                        "specifier");
+        goto done;
+    }
+
+    /* Neither is '=' alignment . */
+    if (format->align == '=') {
+        PyErr_SetString(PyExc_ValueError,
+                        "'=' alignment flag is not allowed in complex format "
+                        "specifier");
+        goto done;
+    }
+
+    re = PyComplex_RealAsDouble(value);
+    if (re == -1.0 && PyErr_Occurred())
+        goto done;
+    im = PyComplex_ImagAsDouble(value);
+    if (im == -1.0 && PyErr_Occurred())
+        goto done;
+
+    if (type == '\0') {
+        /* Omitted type specifier. Should be like str(self). */
+        type = 'g';
+        add_parens = 1;
+        if (re == 0.0)
+            skip_re = 1;
+    }
+
+    if (type == 'n')
+        /* 'n' is the same as 'g', except for the locale used to
+           format the result. We take care of that later. */
+        type = 'g';
+
+    /* 'F' is the same as 'f', per the PEP */
+    if (type == 'F')
+        type = 'f';
+
+    if (precision < 0)
+        precision = 6;
+
+    /* Cast "type", because if we're in unicode we need to pass a
+       8-bit char. This is safe, because we've restricted what "type"
+       can be. */
+    re_buf = PyOS_double_to_string(re, (char)type, precision, flags,
+                                   &re_float_type);
+    if (re_buf == NULL)
+        goto done;
+    im_buf = PyOS_double_to_string(im, (char)type, precision, flags,
+                                   &im_float_type);
+    if (im_buf == NULL)
+        goto done;
+
+    n_re_digits = strlen(re_buf);
+    n_im_digits = strlen(im_buf);
+
+    /* Since there is no unicode version of PyOS_double_to_string,
+       just use the 8 bit version and then convert to unicode. */
+#if STRINGLIB_IS_UNICODE
+    re_unicode_tmp = (Py_UNICODE*)PyMem_Malloc((n_re_digits)*sizeof(Py_UNICODE));
+    if (re_unicode_tmp == NULL) {
+        PyErr_NoMemory();
+        goto done;
+    }
+    strtounicode(re_unicode_tmp, re_buf, n_re_digits);
+    p_re = re_unicode_tmp;
+
+    im_unicode_tmp = (Py_UNICODE*)PyMem_Malloc((n_im_digits)*sizeof(Py_UNICODE));
+    if (im_unicode_tmp == NULL) {
+        PyErr_NoMemory();
+        goto done;
+    }
+    strtounicode(im_unicode_tmp, im_buf, n_im_digits);
+    p_im = im_unicode_tmp;
+#else
+    p_re = re_buf;
+    p_im = im_buf;
+#endif
+
+    /* Is a sign character present in the output?  If so, remember it
+       and skip it */
+    if (*p_re == '-') {
+        re_sign_char = *p_re;
+        ++p_re;
+        --n_re_digits;
+    }
+    if (*p_im == '-') {
+        im_sign_char = *p_im;
+        ++p_im;
+        --n_im_digits;
+    }
+
+    /* Determine if we have any "remainder" (after the digits, might include
+       decimal or exponent or both (or neither)) */
+    parse_number(p_re, n_re_digits, &n_re_remainder, &re_has_decimal);
+    parse_number(p_im, n_im_digits, &n_im_remainder, &im_has_decimal);
+
+    /* Determine the grouping, separator, and decimal point, if any. */
+    get_locale_info(format->type == 'n' ? LT_CURRENT_LOCALE :
+                    (format->thousands_separators ?
+                     LT_DEFAULT_LOCALE :
+                     LT_NO_LOCALE),
+                    &locale);
+
+    /* Turn off any padding. We'll do it later after we've composed
+       the numbers without padding. */
+    tmp_format.fill_char = '\0';
+    tmp_format.align = '\0';
+    tmp_format.width = -1;
+
+    /* Calculate how much memory we'll need. */
+    n_re_total = calc_number_widths(&re_spec, 0, re_sign_char, p_re,
+                                    n_re_digits, n_re_remainder,
+                                    re_has_decimal, &locale, &tmp_format);
+
+    /* Same formatting, but always include a sign. */
+    tmp_format.sign = '+';
+    n_im_total = calc_number_widths(&im_spec, 0, im_sign_char, p_im,
+                                    n_im_digits, n_im_remainder,
+                                    im_has_decimal, &locale, &tmp_format);
+
+    if (skip_re)
+        n_re_total = 0;
+
+    /* Add 1 for the 'j', and optionally 2 for parens. */
+    calc_padding(n_re_total + n_im_total + 1 + add_parens * 2,
+                 format->width, format->align, &lpad, &rpad, &total);
+
+    result = STRINGLIB_NEW(NULL, total);
+    if (result == NULL)
+        goto done;
+
+    /* Populate the memory. First, the padding. */
+    p = fill_padding(STRINGLIB_STR(result),
+                     n_re_total + n_im_total + 1 + add_parens * 2,
+                     format->fill_char=='\0' ? ' ' : format->fill_char,
+                     lpad, rpad);
+
+    if (add_parens)
+        *p++ = '(';
+
+    if (!skip_re) {
+        fill_number(p, &re_spec, p_re, n_re_digits, NULL, 0, &locale, 0);
+        p += n_re_total;
+    }
+    fill_number(p, &im_spec, p_im, n_im_digits, NULL, 0, &locale, 0);
+    p += n_im_total;
+    *p++ = 'j';
+
+    if (add_parens)
+        *p++ = ')';
+
+done:
+    PyMem_Free(re_buf);
+    PyMem_Free(im_buf);
+#if STRINGLIB_IS_UNICODE
+    PyMem_Free(re_unicode_tmp);
+    PyMem_Free(im_unicode_tmp);
+#endif
+    return result;
+}
+#endif /* FORMAT_COMPLEX */
+
+/************************************************************************/
 /*********** built in formatters ****************************************/
 /************************************************************************/
 PyObject *
@@ -1196,3 +1434,50 @@ done:
     return result;
 }
 #endif /* FORMAT_FLOAT */
+
+#ifdef FORMAT_COMPLEX
+PyObject *
+FORMAT_COMPLEX(PyObject *obj,
+               STRINGLIB_CHAR *format_spec,
+               Py_ssize_t format_spec_len)
+{
+    PyObject *result = NULL;
+    InternalFormatSpec format;
+
+    /* check for the special case of zero length format spec, make
+       it equivalent to str(obj) */
+    if (format_spec_len == 0) {
+        result = STRINGLIB_TOSTR(obj);
+        goto done;
+    }
+
+    /* parse the format_spec */
+    if (!parse_internal_render_format_spec(format_spec,
+                                           format_spec_len,
+                                           &format, '\0'))
+        goto done;
+
+    /* type conversion? */
+    switch (format.type) {
+    case '\0': /* No format code: like 'g', but with at least one decimal. */
+    case 'e':
+    case 'E':
+    case 'f':
+    case 'F':
+    case 'g':
+    case 'G':
+    case 'n':
+        /* no conversion, already a complex.  do the formatting */
+        result = format_complex_internal(obj, &format);
+        break;
+
+    default:
+        /* unknown */
+        unknown_presentation_type(format.type, obj->ob_type->tp_name);
+        goto done;
+    }
+
+done:
+    return result;
+}
+#endif /* FORMAT_COMPLEX */
diff --git a/Python/formatter_unicode.c b/Python/formatter_unicode.c
index 206b64037a..c350907da1 100644
--- a/Python/formatter_unicode.c
+++ b/Python/formatter_unicode.c
@@ -6,8 +6,9 @@
 #include "../Objects/stringlib/unicodedefs.h"
 
 
-#define FORMAT_STRING _PyUnicode_FormatAdvanced
-#define FORMAT_LONG   _PyLong_FormatAdvanced
-#define FORMAT_FLOAT  _PyFloat_FormatAdvanced
+#define FORMAT_STRING  _PyUnicode_FormatAdvanced
+#define FORMAT_LONG    _PyLong_FormatAdvanced
+#define FORMAT_FLOAT   _PyFloat_FormatAdvanced
+#define FORMAT_COMPLEX _PyComplex_FormatAdvanced
 
 #include "../Objects/stringlib/formatter.h"