Improve redundant qualifiers removal in generated code

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/CodeGeneratorImpl.kt

@@ -40,6 +40,7 @@ import org.jetbrains.kotlinx.dataframe.keywords.ModifierKeywords
 import org.jetbrains.kotlinx.dataframe.schema.ComparisonMode
 import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema
 import kotlin.reflect.KClass
+import kotlin.text.replace
 
 private fun renderNullability(nullable: Boolean) = if (nullable) "?" else ""
 
@@ -246,18 +247,15 @@ internal object ShortNames : TypeRenderingStrategy {
             }
         }
 
-    private fun String.shorten() = removeRedundantQualifier(this)
-
-    private fun removeRedundantQualifier(markerName: String): String {
-        val parts = markerName.split('.')
-        return if (parts.size == 2 && parts[0] == "kotlin") {
-            parts[1]
-        } else {
-            markerName
-        }
-    }
+    private fun String.shorten() = this.dropKotlinPackages()
 }
 
+internal fun String.dropKotlinPackages() = replace(REDUNDANT_PACKAGE_QUALIFIER, "")
+
+private val REDUNDANT_PACKAGE_QUALIFIER = Regex(
+    """kotlin(?:\.(?:collections|sequences|ranges|text|io|comparisons))?\.(?=[A-Z])""",
+)
+
 internal open class ExtensionsCodeGeneratorImpl(private val typeRendering: TypeRenderingStrategy) :
     ExtensionsCodeGenerator,
     TypeRenderingStrategy by typeRendering {

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt

@@ -543,6 +543,19 @@ class CodeGenerationTests : BaseTest() {
         """^[\d]""".toRegex().matches("3fds")
     }
 
+    @Test
+    fun shortenKotlinCollections() {
+        val df = dataFrameOf("a" to columnOf(listOf("abc")))
+        val expected =
+            """
+            @DataSchema
+            data class DataEntry(
+                val a: List<String>
+            )
+            """.trimIndent()
+        assertEquals(expected, df.generateDataClasses().value)
+    }
+
     // region Tests for generateX functions
 
     @Test

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ShortNamesRenderingTest.kt

@@ -8,6 +8,7 @@ import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
 import org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.ShortNames
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.TypeRenderingStrategy
+import org.jetbrains.kotlinx.dataframe.impl.codeGen.dropKotlinPackages
 import org.junit.Test
 
 internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames {
@@ -61,8 +62,8 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames {
     @Test
     fun `short functional types are not supported`() {
         fields.keys.asClue {
-            fields["d"]!!.renderAccessorFieldType() shouldBe "() -> kotlin.Unit"
-            fields["d"]!!.renderFieldType() shouldBe "() -> kotlin.Unit"
+            fields["d"]!!.renderAccessorFieldType() shouldBe "() -> Unit"
+            fields["d"]!!.renderFieldType() shouldBe "() -> Unit"
         }
     }
 
@@ -112,7 +113,7 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames {
     @Test
     fun `functional type column`() {
         fields.keys.asClue {
-            fields["d"]!!.renderColumnType() shouldBe "DataColumn<() -> kotlin.Unit>"
+            fields["d"]!!.renderColumnType() shouldBe "DataColumn<() -> Unit>"
         }
     }
 
@@ -146,4 +147,14 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames {
     fun `generic column`() {
         MarkersExtractor.get(GenericDataSchema::class).allFields[0].renderColumnType() shouldBe "DataColumn<A>"
     }
+
+    @Test
+    fun `drop redundant kotlin package prefixes`() {
+        "kotlin.String".dropKotlinPackages() shouldBe "String"
+        "kotlin.collections.List<kotlin.String>".dropKotlinPackages() shouldBe "List<String>"
+        "kotlin.collections.Map<kotlin.String, kotlin.collections.List<kotlin.Int>>".dropKotlinPackages() shouldBe
+            "Map<String, List<Int>>"
+        "org.example.Foo<kotlin.String>".dropKotlinPackages() shouldBe "org.example.Foo<String>"
+        "kotlin.internal.Foo".dropKotlinPackages() shouldBe "kotlin.internal.Foo"
+    }
 }