edu.isi.nlp.files.ZipKeyValueSink Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of common-core-open Show documentation
The newest version!
package edu.isi.nlp.files;

import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.base.Charsets;
import com.google.common.base.Function;
import com.google.common.io.Files;
import edu.isi.nlp.symbols.Symbol;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import javax.annotation.Nonnull;

/**
 * A {@link KeyValueSink} based on the contents of a zip file.
 *
 * @author Constantine Lignos, Ryan Gabbard
 */
final class ZipKeyValueSink implements KeyValueSink {

  private final OutputStream output;
  private final ZipOutputStream zip;
  private final Function keyEntryFunction;

  private ZipKeyValueSink(
      final OutputStream outputStream,
      final ZipOutputStream zipOutputStream,
      final Function keyEntryFunction) {
    this.output = checkNotNull(outputStream);
    this.zip = checkNotNull(zipOutputStream);
    this.keyEntryFunction = keyEntryFunction;
  }

  @Nonnull
  static ZipKeyValueSink forFile(
      final File zipfile, final Function keyEntryFunction) throws IOException {
    final OutputStream outputStream = Files.asByteSink(zipfile).openBufferedStream();
    final ZipOutputStream zipOutputStream = new ZipOutputStream(outputStream, Charsets.UTF_8);
    return new ZipKeyValueSink(outputStream, zipOutputStream, keyEntryFunction);
  }

  @Override
  public void put(final Symbol key, final byte[] value) throws IOException {
    final String entryPath = keyEntryFunction.apply(key);
    if (entryPath == null || entryPath.isEmpty()) {
      throw new IllegalArgumentException("Key to entry function returned an invalid string");
    }
    final ZipEntry entry = new ZipEntry(entryPath);
    zip.putNextEntry(entry);
    zip.write(value);
    zip.closeEntry();
  }

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