CsvWriter.java

package org.molgenis.data.csv;

import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import org.molgenis.data.Entity;
import org.molgenis.data.MolgenisDataException;
import org.molgenis.data.convert.DateToStringConverter;
import org.molgenis.data.meta.model.Attribute;
import org.molgenis.data.processor.AbstractCellProcessor;
import org.molgenis.data.processor.CellProcessor;
import org.molgenis.data.support.AbstractWritable;

import java.io.*;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.List;

public class CsvWriter extends AbstractWritable
{
	public static final Charset DEFAULT_CHARSET = Charset.forName("UTF-8");

	public static final char DEFAULT_SEPARATOR = ',';

	private final au.com.bytecode.opencsv.CSVWriter csvWriter;
	/**
	 * process cells before writing
	 */
	private List<CellProcessor> cellProcessors;

	private List<String> cachedAttributeNames;

	public CsvWriter(Writer writer)
	{
		this(writer, ',');
	}

	public CsvWriter(Writer writer, List<String> attributeNames) throws IOException
	{
		this(writer);
		writeAttributeNames(attributeNames);
	}

	public CsvWriter(Writer writer, char separator)
	{
		this(writer, separator, false);
	}

	public CsvWriter(Writer writer, char separator, boolean noQuotes)
	{
		if (writer == null) throw new IllegalArgumentException("writer is null");
		if (noQuotes)
		{
			this.csvWriter = new au.com.bytecode.opencsv.CSVWriter(writer, separator,
					au.com.bytecode.opencsv.CSVWriter.NO_QUOTE_CHARACTER);
		}
		else
		{
			this.csvWriter = new au.com.bytecode.opencsv.CSVWriter(writer, separator);
		}
	}

	public CsvWriter(OutputStream os)
	{
		this(new OutputStreamWriter(os, DEFAULT_CHARSET));
	}

	public CsvWriter(OutputStream os, char separator)
	{
		this(new OutputStreamWriter(os, DEFAULT_CHARSET), separator);
	}

	public CsvWriter(OutputStream os, char separator, boolean noQuotes)
	{
		this(new OutputStreamWriter(os, DEFAULT_CHARSET), separator, noQuotes);
	}

	public CsvWriter(File file) throws FileNotFoundException
	{
		this(new OutputStreamWriter(new FileOutputStream(file), DEFAULT_CHARSET), DEFAULT_SEPARATOR);
	}

	public CsvWriter(File file, char separator) throws FileNotFoundException
	{
		this(new OutputStreamWriter(new FileOutputStream(file), DEFAULT_CHARSET), separator);
	}

	public void addCellProcessor(CellProcessor cellProcessor)
	{
		if (cellProcessors == null) cellProcessors = new ArrayList<CellProcessor>();
		cellProcessors.add(cellProcessor);
	}

	@Override
	public void add(Entity entity)
	{
		if (cachedAttributeNames == null)
			throw new MolgenisDataException("No attribute names defined call writeAttributeNames first");

		int i = 0;
		String[] values = new String[cachedAttributeNames.size()];
		for (String colName : cachedAttributeNames)
		{
			values[i++] = toValue(entity.get(colName));
		}

		csvWriter.writeNext(values);
		if (csvWriter.checkError()) throw new MolgenisDataException("An exception occured writing the csv file");
	}

	public void writeAttributeNames(Iterable<String> attributeNames) throws IOException
	{
		writeAttributes(attributeNames, attributeNames);
	}

	/**
	 * Use attribute labels as column names
	 *
	 * @param attributes
	 * @throws IOException
	 */
	public void writeAttributes(Iterable<Attribute> attributes) throws IOException
	{
		List<String> attributeNames = Lists.newArrayList();
		List<String> attributeLabels = Lists.newArrayList();

		for (Attribute attr : attributes)
		{
			attributeNames.add(attr.getName());
			if (attr.getLabel() != null)
			{
				attributeLabels.add(attr.getLabel());
			}
			else
			{
				attributeLabels.add(attr.getName());
			}
		}

		writeAttributes(attributeNames, attributeLabels);
	}

	public void writeAttributes(Iterable<String> attributeNames, Iterable<String> attributeLabels) throws IOException
	{
		if (cachedAttributeNames == null)
		{
			List<String> processedAttributeNames = new ArrayList<String>();
			for (String colName : attributeNames)
			{
				// process column name
				String processedColName = AbstractCellProcessor.processCell(colName, true, this.cellProcessors);
				processedAttributeNames.add(processedColName);
			}

			// store filtered column names
			cachedAttributeNames = processedAttributeNames;

			// write column labels
			this.csvWriter.writeNext(Iterables.toArray(attributeLabels, String.class));
			if (this.csvWriter.checkError()) throw new IOException();

		}
	}

	@Override
	public void close() throws IOException
	{
		csvWriter.close();
	}

	private String toValue(Object obj)
	{
		String value;
		if (obj == null)
		{
			value = null;
		}
		else if (obj instanceof java.util.Date)
		{
			value = new DateToStringConverter().convert((java.util.Date) obj);
		}
		else if (obj instanceof Entity)
		{
			if (getEntityWriteMode() != null)
			{
				switch (getEntityWriteMode())
				{
					case ENTITY_IDS:
						value = ((Entity) obj).getIdValue().toString();
						break;
					case ENTITY_LABELS:
						Object labelValue = ((Entity) obj).getLabelValue();
						value = labelValue != null ? labelValue.toString() : null;
						break;
					default:
						throw new RuntimeException("Unknown write mode [" + getEntityWriteMode() + "]");
				}
			}
			else
			{
				Object labelValue = ((Entity) obj).getLabelValue();
				value = labelValue != null ? labelValue.toString() : null;
			}
		}
		else if (obj instanceof Iterable<?>)
		{
			StringBuilder strBuilder = new StringBuilder();
			for (Object listItem : (Iterable<?>) obj)
			{
				if (strBuilder.length() > 0) strBuilder.append(',');
				strBuilder.append(toValue(listItem));
			}
			// TODO apply cell processors to list elements?
			value = strBuilder.toString();
		}
		else
		{
			value = obj.toString();
		}
		return AbstractCellProcessor.processCell(value, false, this.cellProcessors);
	}

	@Override
	public void flush()
	{
		try
		{
			csvWriter.flush();
		}
		catch (IOException e)
		{
			throw new MolgenisDataException("Error flushing csvwriter", e);
		}

	}

	@Override
	public void clearCache()
	{
		// Nothing
	}
}