- Added addSheet functionality (keepFile parameter)

 - Updated WorkBook factory as it was previously incompatible with reusing existing workbook
 - Revised test use generated random data

Signed-off-by: LeandroC89 <[email protected]>

Author: LeandroC89

dataframe-excel/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/xlsx.kt

@@ -4,23 +4,12 @@ import kotlinx.datetime.toJavaLocalDate
 import kotlinx.datetime.toJavaLocalDateTime
 import kotlinx.datetime.toKotlinLocalDateTime
 import org.apache.poi.hssf.usermodel.HSSFWorkbook
-import org.apache.poi.ss.usermodel.Cell
-import org.apache.poi.ss.usermodel.CellType
-import org.apache.poi.ss.usermodel.DateUtil
-import org.apache.poi.ss.usermodel.RichTextString
-import org.apache.poi.ss.usermodel.Row
-import org.apache.poi.ss.usermodel.Sheet
-import org.apache.poi.ss.usermodel.Workbook
-import org.apache.poi.ss.usermodel.WorkbookFactory
+import org.apache.poi.ss.usermodel.*
 import org.apache.poi.ss.util.CellReference
 import org.apache.poi.util.LocaleUtil
 import org.apache.poi.util.LocaleUtil.getUserTimeZone
 import org.apache.poi.xssf.usermodel.XSSFWorkbook
-import org.jetbrains.kotlinx.dataframe.AnyFrame
-import org.jetbrains.kotlinx.dataframe.AnyRow
-import org.jetbrains.kotlinx.dataframe.ColumnsSelector
-import org.jetbrains.kotlinx.dataframe.DataColumn
-import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.*
 import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.forEach
 import org.jetbrains.kotlinx.dataframe.api.select
@@ -32,8 +21,7 @@ import java.io.OutputStream
 import java.net.URL
 import java.time.LocalDate
 import java.time.LocalDateTime
-import java.util.Calendar
-import java.util.Date
+import java.util.*
 
 public class Excel : SupportedFormat {
     override fun readDataFrame(stream: InputStream, header: List<String>): AnyFrame = DataFrame.readExcel(stream)
@@ -216,9 +204,10 @@ public fun <T> DataFrame<T>.writeExcel(
     columnsSelector: ColumnsSelector<T, *> = { all() },
     sheetName: String? = null,
     writeHeader: Boolean = true,
-    workBookType: WorkBookType = WorkBookType.XLSX
+    workBookType: WorkBookType = WorkBookType.XLSX,
+    keepFile: Boolean = false,
 ) {
-    return writeExcel(File(path), columnsSelector, sheetName, writeHeader, workBookType)
+    return writeExcel(File(path), columnsSelector, sheetName, writeHeader, workBookType, keepFile)
 }
 
 public enum class WorkBookType {
@@ -230,12 +219,23 @@ public fun <T> DataFrame<T>.writeExcel(
     columnsSelector: ColumnsSelector<T, *> = { all() },
     sheetName: String? = null,
     writeHeader: Boolean = true,
-    workBookType: WorkBookType = WorkBookType.XLSX
+    workBookType: WorkBookType = WorkBookType.XLSX,
+    keepFile: Boolean = false,
 ) {
-    val factory = when (workBookType) {
-        WorkBookType.XLS -> { { HSSFWorkbook() } }
-        WorkBookType.XLSX -> { { XSSFWorkbook() } }
-    }
+
+    val factory =
+        if (keepFile){
+            when (workBookType) {
+                WorkBookType.XLS -> HSSFWorkbook(file.inputStream())
+                WorkBookType.XLSX ->  XSSFWorkbook(file.inputStream())
+            }
+        }
+        else {
+            when (workBookType) {
+                WorkBookType.XLS ->  HSSFWorkbook()
+                WorkBookType.XLSX ->  XSSFWorkbook()
+            }
+        }
     return file.outputStream().use {
         writeExcel(it, columnsSelector, sheetName, writeHeader, factory)
     }
@@ -246,9 +246,9 @@ public fun <T> DataFrame<T>.writeExcel(
     columnsSelector: ColumnsSelector<T, *> = { all() },
     sheetName: String? = null,
     writeHeader: Boolean = true,
-    factory: () -> Workbook
+    factory: Workbook
 ) {
-    val wb: Workbook = factory()
+    val wb: Workbook = factory
     writeExcel(wb, columnsSelector, sheetName, writeHeader)
     wb.write(outputStream)
     wb.close()

dataframe-excel/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/XlsxTest.kt

@@ -111,12 +111,13 @@ class XlsxTest {
     }
 
     @Test
-    fun `write to new sheet when erase is false`() {
-        val df = DataFrame.readExcel(testResource("sample4.xls"), skipRows = 6, rowsCount = 1)
+    fun `write to new sheet when keepFile is true`() {
+        val names = (1..5).map { "column$it" }
+        val df = dataFrameOf(names).randomDouble(7)
         val fileLoc = testResource("generated_wb.xlsx").toURI().toString().removeRange(0, 6)
 
         df.writeExcel(fileLoc, sheetName = "TestSheet1")
-        df.writeExcel(fileLoc, sheetName = "TestSheet2")
+        df.writeExcel(fileLoc, sheetName = "TestSheet2", keepFile = true)
 
         val testSheet1Df = DataFrame.readExcel(testResource("generated_wb.xlsx"), sheetName = "TestSheet1")
         val testSheet2Df = DataFrame.readExcel(testResource("generated_wb.xlsx"), sheetName = "TestSheet2")