pulledIOUtils last updates from origin/dump

hadyelsahar · hadyelsahar · commit 100bcb128be0 · 2013-07-29T17:51:06.000+02:00
diff --git a/core/src/main/scala/org/dbpedia/extraction/util/IOUtils.scala b/core/src/main/scala/org/dbpedia/extraction/util/IOUtils.scala
@@ -1,13 +1,84 @@
 package org.dbpedia.extraction.util
 
-import java.io.{InputStream,OutputStream}
+import org.apache.commons.compress.compressors.bzip2.{BZip2CompressorInputStream,BZip2CompressorOutputStream}
+import java.util.zip.{GZIPInputStream,GZIPOutputStream}
+import java.io.{InputStream,OutputStream,Reader,Writer,OutputStreamWriter,InputStreamReader}
+import scala.io.Codec
+import org.dbpedia.extraction.util.RichReader.wrapReader
+import java.nio.charset.Charset
 
+/**
+ * TODO: modify the bzip code such that there are no run-time dependencies on commons-compress.
+ * Users should be able to use .gz files without having commons-compress on the classpath.
+ * Even better, look for several different bzip2 implementations on the classpath...
+ */
 object IOUtils {
 
+  /**
+   * Map from file suffix (without "." dot) to output stream wrapper
+   */
+  val zippers = Map[String, OutputStream => OutputStream] (
+    "gz" -> { new GZIPOutputStream(_) }, 
+    "bz2" -> { new BZip2CompressorOutputStream(_) } 
+  )
+  
+  /**
+   * Map from file suffix (without "." dot) to input stream wrapper
+   */
+  val unzippers = Map[String, InputStream => InputStream] (
+    "gz" -> { new GZIPInputStream(_) }, 
+    "bz2" -> { new BZip2CompressorInputStream(_, true) } 
+  )
+  
+  /**
+   * use opener on file, wrap in un/zipper stream if necessary
+   */
+  private def open[T](file: FileLike[_], opener: FileLike[_] => T, wrappers: Map[String, T => T]): T = {
+    val name = file.name
+    val suffix = name.substring(name.lastIndexOf('.') + 1)
+    wrappers.getOrElse(suffix, identity[T] _)(opener(file)) 
+  }
+  
+  /**
+   * open output stream, wrap in zipper stream if file suffix indicates compressed file.
+   */
+  def outputStream(file: FileLike[_]): OutputStream =
+    open(file, _.outputStream(), zippers)
+  
+  /**
+   * open input stream, wrap in unzipper stream if file suffix indicates compressed file.
+   */
+  def inputStream(file: FileLike[_]): InputStream =
+    open(file, _.inputStream(), unzippers)
+  
+  /**
+   * open output stream, wrap in zipper stream if file suffix indicates compressed file,
+   * wrap in writer.
+   */
+  def writer(file: FileLike[_], charset: Charset = Codec.UTF8.charSet): Writer =
+    new OutputStreamWriter(outputStream(file), charset)
+  
+  /**
+   * open input stream, wrap in unzipper stream if file suffix indicates compressed file,
+   * wrap in reader.
+   */
+  def reader(file: FileLike[_], charset: Charset = Codec.UTF8.charSet): Reader =
+    new InputStreamReader(inputStream(file), charset)
+  
+  def readLines(file: FileLike[_])(proc: String => Unit): Unit = {
+    val reader = this.reader(file)
+    try {
+      for (line <- reader) {
+        proc(line)
+      }
+    }
+    finally reader.close()
+  }
+
   /**
    * Copy all bytes from input to output. Don't close any stream.
    */
-  def copy(in : InputStream, out : OutputStream) : Unit = {
+  def copy(in: InputStream, out: OutputStream) : Unit = {
     val buf = new Array[Byte](1 << 20) // 1 MB
     while (true)
     {
@@ -21,4 +92,4 @@ object IOUtils {
     }
   }
 
-}
+}
