objcomplex.c

#include <stdlib.h>
#include <assert.h>

#include "nlr.h"
#include "misc.h"
#include "mpconfig.h"
#include "qstr.h"
#include "obj.h"
#include "parsenum.h"
#include "runtime0.h"

#if MICROPY_ENABLE_FLOAT

#if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
#include "formatfloat.h"
#endif

typedef struct _mp_obj_complex_t {
    mp_obj_base_t base;
    mp_float_t real;
    mp_float_t imag;
} mp_obj_complex_t;

mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag);

STATIC void complex_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t o_in, mp_print_kind_t kind) {
    mp_obj_complex_t *o = o_in;
#if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
    char buf[32];
    if (o->real == 0) {
        format_float(o->imag, buf, sizeof(buf), 'g', 6, '\0');
        print(env, "%sj", buf);
    } else {
        format_float(o->real, buf, sizeof(buf), 'g', 6, '\0');
        print(env, "(%s+", buf);
        format_float(o->imag, buf, sizeof(buf), 'g', 6, '\0');
        print(env, "%sj)", buf);
    }
#else
    if (o->real == 0) {
        print(env, "%.8gj",  (double) o->imag);
    } else {
        print(env, "(%.8g+%.8gj)", (double) o->real, (double) o->imag);
    }
#endif
}

STATIC mp_obj_t complex_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
    // TODO check n_kw == 0

    switch (n_args) {
        case 0:
            return mp_obj_new_complex(0, 0);

        case 1:
            if (MP_OBJ_IS_STR(args[0])) {
                // a string, parse it
                uint l;
                const char *s = mp_obj_str_get_data(args[0], &l);
                return mp_parse_num_decimal(s, l, true, true);
            } else if (MP_OBJ_IS_TYPE(args[0], &mp_type_complex)) {
                // a complex, just return it
                return args[0];
            } else {
                // something else, try to cast it to a complex
                return mp_obj_new_complex(mp_obj_get_float(args[0]), 0);
            }

        case 2: {
            mp_float_t real, imag;