Merge pull request #107 from Kotlin/fix-csv-write

Fix csv write

Author: koperagen

plugins/dataframe-gradle-plugin/src/integrationTest/kotlin/org/jetbrains/dataframe/gradle/SchemaGeneratorPluginIntegrationTest.kt

@@ -218,7 +218,8 @@ class SchemaGeneratorPluginIntegrationTest : AbstractDataFramePluginIntegrationT
         fun escapeDoubleQuotes(it: Char) = if (it == '"') "\"\"" else it.toString()
 
         val (_, result) = runGradleBuild(":build") { buildDir ->
-            val dataFile = File(buildDir, "data.csv")
+            val filename = "data.csv"
+            val dataFile = File(buildDir, filename)
             val notSupportedChars = setOf('\n', '\r')
             (Char.MIN_VALUE..Char.MAX_VALUE).asSequence()
                 .filterNot { it in notSupportedChars }
@@ -239,7 +240,7 @@ class SchemaGeneratorPluginIntegrationTest : AbstractDataFramePluginIntegrationT
                 import org.jetbrains.kotlinx.dataframe.api.filter
                 
                 fun main() {
-                    val df = DataFrame.read("$dataFile").cast<Schema>()
+                        val df = DataFrame.read("${TestData.csvName}").cast<Schema>()
                 }
             """.trimIndent())
 
@@ -264,7 +265,7 @@ class SchemaGeneratorPluginIntegrationTest : AbstractDataFramePluginIntegrationT
 
                 dataframes {
                     schema {
-                        data = "$dataFile"
+                        data = "${TestData.csvName}"
                         name = "Schema"
                         packageName = ""
                     }
@@ -295,7 +296,7 @@ class SchemaGeneratorPluginIntegrationTest : AbstractDataFramePluginIntegrationT
                 }
                 
                 fun main() {
-                    val df = DataFrame.read("$dataFile").cast<MySchema>()
+                    val df = DataFrame.read("${TestData.csvName}").cast<MySchema>()
                     val df1 = df.filter { age != null }
                 }
             """.trimIndent())

src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/csv.kt

@@ -6,10 +6,7 @@ import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
-import org.jetbrains.kotlinx.dataframe.api.ParserOptions
-import org.jetbrains.kotlinx.dataframe.api.forEachRow
-import org.jetbrains.kotlinx.dataframe.api.toDataFrame
-import org.jetbrains.kotlinx.dataframe.api.tryParse
+import org.jetbrains.kotlinx.dataframe.api.*
 import org.jetbrains.kotlinx.dataframe.impl.ColumnNameGenerator
 import org.jetbrains.kotlinx.dataframe.impl.api.Parsers
 import org.jetbrains.kotlinx.dataframe.impl.api.parse
@@ -296,13 +293,22 @@ public fun DataFrame.Companion.readDelim(
     return cols.toDataFrame()
 }
 
-public fun AnyFrame.writeCSV(file: File, format: CSVFormat = CSVFormat.DEFAULT.withHeader()): Unit =
+public fun AnyFrame.writeCSV(
+    file: File,
+    format: CSVFormat = CSVFormat.DEFAULT
+): Unit =
     writeCSV(FileWriter(file), format)
 
-public fun AnyFrame.writeCSV(path: String, format: CSVFormat = CSVFormat.DEFAULT.withHeader()): Unit =
+public fun AnyFrame.writeCSV(
+    path: String,
+    format: CSVFormat = CSVFormat.DEFAULT
+): Unit =
     writeCSV(FileWriter(path), format)
 
-public fun AnyFrame.writeCSV(writer: Appendable, format: CSVFormat = CSVFormat.DEFAULT.withHeader()): Unit =
+public fun AnyFrame.writeCSV(
+    writer: Appendable,
+    format: CSVFormat = CSVFormat.DEFAULT
+) {
     format.print(writer).use { printer ->
         printer.printRecord(columnNames())
         forEachRow {
@@ -316,8 +322,9 @@ public fun AnyFrame.writeCSV(writer: Appendable, format: CSVFormat = CSVFormat.D
             printer.printRecord(values)
         }
     }
+}
 
-public fun AnyFrame.toCsv(format: CSVFormat = CSVFormat.DEFAULT.withHeader()): String =
+public fun AnyFrame.toCsv(format: CSVFormat = CSVFormat.DEFAULT): String =
     StringWriter().use {
         this.writeCSV(it, format)
         it

src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/CsvTests.kt

@@ -180,6 +180,12 @@ class CsvTests {
         res shouldBe grouped
     }
 
+    @Test
+    fun `CSV String of saved dataframe starts with column name`() {
+        val df = dataFrameOf("a")(1)
+        df.toCsv().first() shouldBe 'a'
+    }
+
     @Test
     fun `guess tsv`() {
         val df = DataFrame.read(testResource("abc.tsv"))