From e5d798b86742797d95b82ed6ee0301e9d481ba14 Mon Sep 17 00:00:00 2001
From: 0xnm <0xnm@users.noreply.github.com>
Date: Sat, 28 Dec 2024 22:20:48 +0100
Subject: [PATCH] Specify exact files to process in ktlint run

---
 .../mill/kotlinlib/ktlint/KtlintModule.scala  | 34 ++++++++++++-------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/kotlinlib/src/mill/kotlinlib/ktlint/KtlintModule.scala b/kotlinlib/src/mill/kotlinlib/ktlint/KtlintModule.scala
index 8c6a43b3315..95f5fc759e8 100644
--- a/kotlinlib/src/mill/kotlinlib/ktlint/KtlintModule.scala
+++ b/kotlinlib/src/mill/kotlinlib/ktlint/KtlintModule.scala
@@ -1,10 +1,12 @@
 package mill.kotlinlib.ktlint
 
+import mainargs.arg
 import mill._
 import mill.api.{Loose, PathRef}
 import mill.define.{Discover, ExternalModule}
 import mill.javalib.JavaModule
 import mill.kotlinlib.DepSyntax
+import mill.main.Tasks
 import mill.util.Jvm
 
 /**
@@ -16,7 +18,13 @@ trait KtlintModule extends JavaModule {
    * Runs [[https://pinterest.github.io/ktlint/latest/install/integrations/ Ktlint]]
    */
   def ktlint(@mainargs.arg ktlintArgs: KtlintArgs): Command[Unit] = Task.Command {
-    KtlintModule.ktlintAction(ktlintArgs, ktlintConfig(), ktlintOptions(), ktlintClasspath())
+    KtlintModule.ktlintAction(
+      ktlintArgs,
+      sources(),
+      ktlintConfig(),
+      ktlintOptions(),
+      ktlintClasspath()
+    )
   }
 
   /**
@@ -57,15 +65,14 @@ object KtlintModule extends ExternalModule with KtlintModule with TaskModule {
 
   /**
    * Reformats Kotlin source files.
-   *
-   * @param check if an exception should be raised when formatting errors are found
-   *              - when set, files are not formatted
    */
   def reformatAll(
-      check: mainargs.Flag = mainargs.Flag(value = false)
+      @arg(positional = true) sources: Tasks[Seq[PathRef]] =
+        Tasks.resolveMainDefault("__.sources")
   ): Command[Unit] = Task.Command {
     ktlintAction(
-      KtlintArgs(format = true, check = check.value),
+      KtlintArgs(format = true, check = true),
+      T.sequence(sources.value)().flatten,
       ktlintConfig(),
       ktlintOptions(),
       ktlintClasspath()
@@ -74,15 +81,14 @@ object KtlintModule extends ExternalModule with KtlintModule with TaskModule {
 
   /**
    * Checks the Kotlin source files formatting without reformatting the files.
-   *
-   * @param check if an exception should be raised when formatting errors are found
-   *              - when set, files are not formatted
    */
   def checkFormatAll(
-      check: mainargs.Flag = mainargs.Flag(value = false)
+      @arg(positional = true) sources: Tasks[Seq[PathRef]] =
+        Tasks.resolveMainDefault("__.sources")
   ): Command[Unit] = Task.Command {
     ktlintAction(
-      KtlintArgs(format = false, check = check.value),
+      KtlintArgs(format = false, check = false),
+      T.sequence(sources.value)().flatten,
       ktlintConfig(),
       ktlintOptions(),
       ktlintClasspath()
@@ -91,6 +97,7 @@ object KtlintModule extends ExternalModule with KtlintModule with TaskModule {
 
   private def ktlintAction(
       ktlintArgs: KtlintArgs,
+      filesToFormat: Seq[PathRef],
       config: Option[PathRef],
       options: Seq[String],
       classPath: Loose.Agg[PathRef]
@@ -111,12 +118,15 @@ object KtlintModule extends ExternalModule with KtlintModule with TaskModule {
     args ++= options
     args ++= configArgument
     args ++= formatArgument
+    args ++= filesToFormat.map(_.path)
+      .filter(f => os.exists(f) && (f.ext == "kt" || f.ext == "kts"))
+      .map(_.toString())
 
     val exitCode = Jvm.callSubprocess(
       mainClass = "com.pinterest.ktlint.Main",
       classPath = classPath.map(_.path),
       mainArgs = args.result(),
-      workingDir = millSourcePath, // allow passing relative paths for sources like src/a/b
+      workingDir = millSourcePath,
       streamOut = true,
       check = false
     ).exitCode