fucusy · fucusy · Apr 20, 2020 · Apr 18, 2020 · Apr 20, 2020
diff --git a/build.sbt b/build.sbt
@@ -1,6 +1,6 @@
 organization := "fucusy"
 name := "dataframe2html"
-version := "0.1.1-SNAPSHOT"
+version := "0.1.2-SNAPSHOT"
 scalaVersion := "2.11.10"
 
 val sparkVersion  = "2.4.3"

diff --git a/src/main/scala/io/github/fucusy/Viz.scala b/src/main/scala/io/github/fucusy/Viz.scala
@@ -3,7 +3,6 @@ package io.github.fucusy
 import org.apache.spark.sql.{DataFrame, Row}
 
 object Viz {
-
   def imgUrl2tag(url: String) = s"""<img src="$url"/>"""
 
   def isImgUrl(s: String): Boolean = {
@@ -25,20 +24,39 @@ object Viz {
    * B -> Seq[....],
    * C -> Seq[....],
    * D -> Seq[....]) the seq value must in order
-   * - imageCol: tell us which column is image.
    * - columnNames
    *
    * second step: using the four information got form step 1 to generate html
    *
    * @param df
-   * @param imgCols
-   * @param title : the description of whole df
+   * @param imgCols : tell us which column is image.
+   * @param title   : the description of whole df
    * @param limitShowNumber
    */
   def dataframe2html(df: DataFrame,
                      imgCols: Seq[String],
                      title: String,
                      limitShowNumber: Int = -1): String = {
+    val contentInfo = dataframe2data(df, limitShowNumber)
+    val html = data2html(contentInfo, imgCols, title)
+    html
+  }
+
+  /**
+   * Automatically convert img url to img tag
+   * @param df
+   * @param title   : the description of whole df
+   * @param limitShowNumber
+   */
+  def dataframe2html(df: DataFrame,
+                     title: String,
+                     limitShowNumber: Int): String = {
+    val contentInfo = dataframe2data(df, limitShowNumber)
+    val html = data2html(contentInfo, title)
+    html
+  }
+
+  def dataframe2data(df: DataFrame, limitShowNumber: Int = -1): Seq[(String, Seq[String])] = {
     val columnNames: Seq[String] = df.columns
     val collectDF = df.select(columnNames.head, columnNames.tail: _*)
       .collect()
@@ -49,14 +67,12 @@ object Viz {
     } else {
       limitShowNumber
     }
-    val contentInfo = columnNames.zipWithIndex.map {
+    columnNames.zipWithIndex.map {
       case (col: String, idx: Int) => (col, collectDF.map(_ (idx)).slice(0, actualLimitShowNumber).toSeq)
     }
-
-    val html = data2html(contentInfo, imgCols, title)
-    html
   }
 
+
   /** *
    * convert data to html
    *

diff --git a/src/main/scala/io/github/fucusy/VizImplicit.scala b/src/main/scala/io/github/fucusy/VizImplicit.scala
@@ -0,0 +1,12 @@
+package io.github.fucusy
+
+import org.apache.spark.sql.DataFrame
+
+object VizImplicit {
+
+  implicit class VizDataFrame(df: DataFrame) {
+    def toHTML(title: String = "", limit: Int = 100): String = {
+      Viz.dataframe2html(df, title, limit)
+    }
+  }
+}
diff --git a/src/test/resources/data2html_no_title.html b/src/test/resources/data2html_no_title.html
@@ -0,0 +1,15 @@
+<html>
+<body>
+<h3></h3>
+<table>
+    <tbody>
+    <tr><th>name</th><td>bob</td><td>amy</td></tr>
+    <tr><th>hobby</th><td>dance</td><td>swimming</td></tr>
+    <tr><th>picture</th>
+        <td><img src="https://raw.githubusercontent.com/fucusy/dataframe2html/57c8b41dfc7368ad7d371c0a94614412abfb1de6/src/test/resources/bob.png"/></td>
+        <td><img src="https://raw.githubusercontent.com/fucusy/dataframe2html/57c8b41dfc7368ad7d371c0a94614412abfb1de6/src/test/resources/amy.png"/></td>
+    </tr>
+    </tbody>
+</table>
+</body>
+</html>
diff --git a/src/test/scala/io/github/fucusy/TestVizImplicit.scala b/src/test/scala/io/github/fucusy/TestVizImplicit.scala
@@ -0,0 +1,28 @@
+package io.github.fucusy
+
+import io.github.fucusy.VizImplicit.VizDataFrame
+import org.apache.spark._
+import org.apache.spark.sql.SparkSession
+import org.scalatest.funsuite.AnyFunSuite
+
+import scala.io.Source
+
+class TestVizImplicit extends AnyFunSuite {
+  test("Test dataframe2html auto convert img url") {
+
+    val sparkConf = new SparkConf()
+      .setMaster("local[*]")
+      .setAppName("hello world")
+    lazy val spark: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
+    import spark.implicits._
+    val c2d = Seq(
+      ("bob", "dance", "https://raw.githubusercontent.com/fucusy/dataframe2html/57c8b41dfc7368ad7d371c0a94614412abfb1de6/src/test/resources/bob.png"),
+      ("amy", "swimming", "https://raw.githubusercontent.com/fucusy/dataframe2html/57c8b41dfc7368ad7d371c0a94614412abfb1de6/src/test/resources/amy.png")
+    ).toDF("name", "hobby", "picture")
+    val title = "Users"
+    val html = c2d.toHTML()
+    val trueHtml = Source.fromFile("src/test/resources/data2html_no_title.html").mkString("")
+    assert(html.replaceAll("\\s", "") == trueHtml.replaceAll("\\s", "")
+    )
+  }
+}