PyWriter.java

/* Copyright (c) Jython Developers */
package org.python.modules._csv;

import org.python.core.Py;
import org.python.core.PyException;
import org.python.core.PyObject;
import org.python.core.PyString;
import org.python.core.PyType;
import org.python.expose.ExposedType;
import org.python.expose.ExposedMethod;
import org.python.expose.ExposedGet;

/**
 * CSV file writer.
 *
 * Analogous to CPython's _csv.c::WriterObj struct.
 */
@ExposedType(name = "_csv.writer")
public class PyWriter extends PyObject {

    public static final PyType TYPE = PyType.fromClass(PyWriter.class);

    /** Parsing dialect. */
    @ExposedGet
    public PyDialect dialect;

    /** Output lines writer callable. */
    private PyObject writeline;

    /** Buffer for parser.join. */
    private StringBuffer rec;

    /** Length of record. */
    private int rec_len = 0;

    /** Number of fields in record. */
    private int num_fields = 0;

    /** Whether field should be quoted during a join. */
    private boolean quoted = false;

    public PyWriter(PyObject writeline, PyDialect dialect) {
        this.writeline = writeline;
        this.dialect = dialect;
    }

    public static PyString __doc__writerows = Py.newString(
            "writerows(sequence of sequences)\n" +
            "\n" +
            "Construct and write a series of sequences to a csv file.  Non-string\n" +
            "elements will be converted to string.");

    public void writerows(PyObject seqseq) {
        writer_writerows(seqseq);
    }

    @ExposedMethod
    final void writer_writerows(PyObject seqseq) {
        PyObject row_iter;
        PyObject row_obj;
        boolean result;

        row_iter = seqseq.__iter__();
        if (row_iter == null) {
            throw _csv.Error("writerows() argument must be iterable");
        }

        while ((row_obj = row_iter.__iternext__()) != null) {
            result = writerow(row_obj);
            if (!result) {
                break;
            }
        }
    }

    public static PyString __doc__writerow = Py.newString(
            "writerow(sequence)\n" +
            "\n" +
            "Construct and write a CSV record from a sequence of fields.  Non-string\n" +
            "elements will be converted to string."
            );

    public boolean writerow(PyObject seq) {
        return writer_writerow(seq);
    }

    @ExposedMethod
    final boolean writer_writerow(PyObject seq) {
        int len;
        int i;

        if (!seq.isSequenceType()) {
            throw _csv.Error("sequence expected");
        }

        len = seq.__len__();
        if (len < 0) {
            return false;
        }

        // Join all fields in internal buffer.
        join_reset();
        for (i = 0; i < len; i++) {
            PyObject field;
            boolean append_ok;
            quoted = false;

            field = seq.__getitem__(i);
            if (field == null) {
                return false;
            }

            switch (dialect.quoting) {
                case QUOTE_NONNUMERIC:
                    try {
                        field.__float__();
                    } catch (PyException ex) {
                        quoted = true;
                    }
                    break;
                case QUOTE_ALL:
                    quoted = true;
                    break;
                default:
                    quoted = false;
            }

            if (field instanceof PyString) {
                append_ok = join_append(field.toString(), len == 1);
            } else if (field == Py.None) {
                append_ok = join_append("", len == 1);
            } else {
                PyObject str = field.__str__();
                if (str == null) {
                    return false;
                }

                append_ok = join_append(str.toString(), len == 1);
            }
            if (!append_ok) {
                return false;
            }
        }

        // Add line terminator.
        if (!join_append_lineterminator()) {
            return false;
        }

        writeline.__call__(new PyString(rec.toString()));
        return true;
    }

    private void join_reset() {
        rec_len = 0;
        num_fields = 0;
        quoted = false;
        rec = new StringBuffer();
    }

    private boolean join_append_lineterminator() {
        rec.append(dialect.lineterminator);
        return true;
    }

    private boolean join_append(String field, boolean quote_empty) {
        int rec_len;

        rec_len = join_append_data(field, quote_empty, false);
        if (rec_len < 0) {
            return false;
        }

        this.rec_len = join_append_data(field, quote_empty, true);
        num_fields++;

        return true;
    }

    /**
     * This method behaves differently depending on the value of copy_phase: if copy_phase
     * is false, then the method determines the new record length. If copy_phase is true
     * then the new field is appended to the record.
     */
    private int join_append_data(String field, boolean quote_empty, boolean copy_phase) {
        int i;

        // If this is not the first field we need a field separator.
        if (num_fields > 0) {
            addChar(dialect.delimiter, copy_phase);
        }

        // Handle preceding quote
        if (copy_phase && quoted) {
            addChar(dialect.quotechar, copy_phase);
        }

        // parsing below is based on _csv.c which expects all strings to be terminated
        // with a nul byte.
        field += '\0';

        // Copy/count field data.
        for (i = 0;; i++) {
            char c = field.charAt(i);
            boolean want_escape = false;

            if (c == '\0') {
                break;
            }
            if (c == dialect.delimiter || c == dialect.escapechar || c == dialect.quotechar
                || dialect.lineterminator.indexOf(c) > -1) {
                if (dialect.quoting == QuoteStyle.QUOTE_NONE) {
                    want_escape = true;
                } else {
                    if (c == dialect.quotechar) {
                        if (dialect.doublequote) {
                            addChar(dialect.quotechar, copy_phase);
                        } else {
                            want_escape = true;
                        }
                    }
                    if (!want_escape) {
                        quoted = true;
                    }
                }
                if (want_escape) {
                    if (dialect.escapechar == '\0') {
                        throw _csv.Error("need to escape, but no escapechar set");
                    }
                    addChar(dialect.escapechar, copy_phase);
                }
            }

            // Copy field character into record buffer.
            addChar(c, copy_phase);
        }

        // If field is empty check if it needs to be quoted.
        if (i == 0 && quote_empty) {
            if (dialect.quoting == QuoteStyle.QUOTE_NONE) {
                throw _csv.Error("single empty field record must be quoted");
            } else {
                quoted = true;
            }
        }

        // Handle final quote character on field.
        if (quoted) {
            if (copy_phase) {
                addChar(dialect.quotechar, copy_phase);
            } else {
                // Didn't know about leading quote until we found it necessary in field
                // data - compensate for it now.
                rec_len += 2;
            }
        }

        return rec_len;
    }

    private void addChar(char c, boolean copy_phase) {
        if (copy_phase) {
            rec.append(c);
        }
        rec_len++;
    }
}