net.sf.okapi.lib.serialization.writer.ProtoBufferTextUnitFlatWriter Maven / Gradle / Ivy

Go to download
/*===========================================================================
  Copyright (C) 2016 by the Okapi Framework contributors
-----------------------------------------------------------------------------
  This library 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 2.1 of the License, or (at 
  your option) any later version.

  This library 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 Lesser 
  General Public License for more details.

  You should have received a copy of the GNU Lesser General Public License 
  along with this library; if not, write to the Free Software Foundation, 
  Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

  See also the full LGPL text here: http://www.gnu.org/copyleft/lesser.html
===========================================================================*/

package net.sf.okapi.lib.serialization.writer;

import com.google.protobuf.TextFormat;
import com.google.protobuf.util.JsonFormat;
import net.sf.okapi.common.Event;
import net.sf.okapi.common.EventType;
import net.sf.okapi.common.IParameters;
import net.sf.okapi.common.LocaleId;
import net.sf.okapi.common.encoder.EncoderManager;
import net.sf.okapi.common.exceptions.OkapiIOException;
import net.sf.okapi.common.filterwriter.IFilterWriter;
import net.sf.okapi.common.resource.ITextUnit;
import net.sf.okapi.common.resource.TextContainer;
import net.sf.okapi.common.skeleton.ISkeletonWriter;
import net.sf.okapi.lib.serialization.textunitflat.TextUnitFlat2Proto;
import net.sf.okapi.proto.textunitflat.TextUnits;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.nio.charset.StandardCharsets;
import java.util.LinkedList;
import java.util.List;

public class ProtoBufferTextUnitFlatWriter implements IFilterWriter {
	private OutputStream outputStream;
	private String outputPath;
	private LocaleId sourceLocale;
	private LocaleId targetLocale;
	private boolean canceled;
	private Parameters params;
	private OutputStreamWriter outputWriter;
	private final List textUnits;

	public ProtoBufferTextUnitFlatWriter() {
		this.params = new Parameters();
		textUnits = new LinkedList<>();
	}
	
	@Override
	public String getName() {
		return getClass().getName();
	}

	@Override
	public void setOptions(LocaleId targetLocale, String defaultEncoding) {
		this.targetLocale = targetLocale;
		// defaultEncoding is ignored: we always use UTF-8
	}

	@Override
	public void setOutput(String path) {
		this.outputPath = path;
	}

	@Override
	public void setOutput(OutputStream output) {
		// If we use the stream, we can't use the path
		this.outputPath = null;
		this.outputStream = output;
	}

	@Override
	public Event handleEvent(Event event) {
		if (canceled) {
			return new Event(EventType.CANCELED);
		}
		switch (event.getEventType()) {
		case START_DOCUMENT:
			processStartDocument(event);
			break;
		case TEXT_UNIT:
			processTextUnit(event);
			break;
		case END_DOCUMENT:
			saveTextUnits();
			break;
		default:
			event = Event.createNoopEvent();
			break;
		}
		return event;
	}

	private void processTextUnit(Event event) {
		if (params.isCopySource() && (!event.getTextUnit().hasTarget(targetLocale) || event.getTextUnit().getTarget(targetLocale).isEmpty())) {
			TextContainer tc = event.getTextUnit().getSource().clone();
			event.getTextUnit().setTarget(targetLocale, tc);
		}
		textUnits.add(event.getTextUnit());
	}

	private void saveTextUnits() {
		switch (params.getOutput()) {
			case TEXT:
				saveToTextFile(outputStream,
						TextUnitFlat2Proto.toTextUnits(textUnits, sourceLocale));
				break;
			case JSON:
				saveToJsonFile(outputStream,
						TextUnitFlat2Proto.toTextUnits(textUnits, sourceLocale));
				break;
			case BINARY:
				saveToBinaryFile(outputStream,
						TextUnitFlat2Proto.toTextUnits(textUnits, sourceLocale));
				break;
		}
	}

	private void saveToTextFile(OutputStream outStream, TextUnits proto) {
		try {
			final String text = TextFormat.printer().escapingNonAscii(false).printToString(proto);
			outStream.write(text.getBytes(StandardCharsets.UTF_8));
		} catch (IOException e) {
			throw new OkapiIOException(e);
		}
	}

	private void saveToJsonFile(OutputStream outStream, TextUnits proto) {
		try {
			final String jsonText = JsonFormat.printer().print(proto);
			outStream.write(jsonText.getBytes(StandardCharsets.UTF_8));
		} catch (IOException e) {
			throw new OkapiIOException(e);
		}
	}

	private void saveToBinaryFile(OutputStream outStream, TextUnits proto) {
		try {
			proto.writeTo(outStream);
		} catch (IOException e) {
			throw new OkapiIOException(e);
		}
	}
	
	private void processStartDocument(Event event) {
		if ( outputStream == null ) {	
			try {
				outputStream = new FileOutputStream(outputPath);
			} catch (FileNotFoundException e) {
				throw new OkapiIOException("Error serializing TextUnit", e);
			}
		}
		
		sourceLocale = event.getStartDocument().getLocale();
	}

	@Override
	public void close() {
		if (outputWriter != null) {
			try {
				outputWriter.close();
				outputWriter = null;
				outputStream = null;
			} catch (IOException e) {
				throw new OkapiIOException("Error closing JSON writer", e);
			}
		}
		
		if (outputStream != null) {
			try {
				outputStream.close();
				outputStream = null;
			} catch (IOException e) {
				throw new OkapiIOException("Error closing protobuf stream", e);
			}
		}
	}

	@Override
	public Parameters getParameters() {
		return params;
	}

	@Override
	public void setParameters(IParameters params) {
		this.params = (Parameters)params;		
	}

	@Override
	public void cancel() {
		close();
		this.canceled = true;
	}

	@Override
	public EncoderManager getEncoderManager() {
		return null;
	}

	@Override
	public ISkeletonWriter getSkeletonWriter() {
		return null;
	}
}