/* This program is free software: you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public License
as published by the Free Software Foundation, either version 3 of
the License, or (at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>. */
package org.opentripplanner.scripting.api;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import com.csvreader.CsvWriter;
import com.google.common.base.Charsets;
/**
* This class allow one to generate easily tabular data and save it as a CSV file.
*
* For example, in python:
* <pre>
* csv = otp.createCSVOutput()
* csv.setHeader( [ 'lat', 'lon', 'total' ] )
* csv.addRow( [ 45.123, 5.789, 42 ] )
* csv.addRow( [ 45.124, 5.792, 34 ] )
* csv.save('mydata.csv')
* </pre>
*
* TODO Rename this class to "TabularOutput" and allow for saving in various format.
*
* @author laurent
*/
public class OtpsCsvOutput {
private List<List<String>> data = new ArrayList<>();
private String[] headers;
protected OtpsCsvOutput() {
}
/**
* Set the (optional) column header names. If this is not set, no header will be generated.
*
* @param headers An array of string, each entry is the name of the corresponding column header,
* in order.
*/
public void setHeader(Object[] headers) {
this.headers = new String[headers.length];
for (int i = 0; i < headers.length; i++) {
this.headers[i] = headers[i].toString();
}
}
/**
* Add a new row to the data.
*
* @param row An array of objects. The order and size of the array should correspond to the
* header. The default toString method of each object will be called to get the actual
* data to output; so any type of object can be provided (string, numbers...)
*/
public void addRow(Object[] row) {
List<String> strs = new ArrayList<>(row.length);
for (int i = 0; i < row.length; i++) {
strs.add(row[i] == null ? "" : row[i].toString());
}
data.add(strs);
}
/**
* Save the data to a file.
* @param file The name of the file to save the data to.
* @throws IOException In case something bad happens (IO exception)
*/
public void save(String file) throws IOException {
CsvWriter csvWriter = new CsvWriter(file, ',', Charsets.UTF_8);
if (headers != null)
csvWriter.writeRecord(headers);
for (List<String> row : data) {
csvWriter.writeRecord(row.toArray(new String[row.size()]));
}
csvWriter.close();
}
/**
* @return The CSV data as a string. It can be used for example as the script return value.
* @see OtpsEntryPoint.setRetval()
* @throws IOException
*/
public String asText() throws IOException {
ByteArrayOutputStream baos = new ByteArrayOutputStream();
CsvWriter csvWriter = new CsvWriter(baos, ',', Charsets.UTF_8);
if (headers != null)
csvWriter.writeRecord(headers);
for (List<String> row : data) {
csvWriter.writeRecord(row.toArray(new String[row.size()]));
}
csvWriter.close();
return new String(baos.toByteArray(), Charsets.UTF_8);
}
}