Some improvements to pythonlib

example/pythonlib/basic/1-simple/build.mill

@@ -9,8 +9,15 @@ object foo extends PythonModule {
   def pythonDeps = Seq("numpy==2.1.3")
 }
 
-object qux extends PythonModule {
+object qux extends PythonModule { q =>
   def moduleDeps = Seq(foo, foo.bar)
+
+  object test extends PythonTests with TestModule.Unittest
+  object test2 extends PythonTests with TestModule.Pytest {
+    override def sources = T{
+      q.test.sources()
+    }
+  }
 }
 
 /** Usage

example/pythonlib/basic/1-simple/qux/src/main.py

@@ -1,11 +1,11 @@
 #!/usr/bin/python3
 import numpy as np
 
-from foo.src.foo import data
-from foo.bar.src.bar import df
+from foo import data
+from bar import df
 
 def main() -> None:
     print(f"Numpy : Sum: {np.sum(data)} | Pandas: Mean: {df['Values'].mean()}, Max: {df['Values'].max()}")
 
 if __name__ == "__main__":
-    main()
+    main()

example/pythonlib/basic/1-simple/qux/test/src/test_dummt.py

@@ -0,0 +1,17 @@
+import unittest
+
+class TestStringMethods(unittest.TestCase):
+
+    def test_upper(self):
+        self.assertEqual('foo'.upper(), 'FOO')
+
+    def test_isupper(self):
+        self.assertTrue('FOO'.isupper())
+        self.assertFalse('Foo'.isupper())
+
+    def test_split(self):
+        s = 'hello world'
+        self.assertEqual(s.split(), ['hello', 'world'])
+        # check that s.split fails when the separator is not a string
+        with self.assertRaises(TypeError):
+            s.split(2)

pythonlib/src/mill/pythonlib/PythonModule.scala

@@ -1,10 +1,24 @@
 package mill.pythonlib
+
 import mill._
+import mill.api.Result
+import mill.util.Util
+import mill.util.Jvm
 
-trait PythonModule extends Module {
+trait PythonModule extends Module with TaskModule { outer =>
   def moduleDeps: Seq[PythonModule] = Nil
-  def mainFileName: T[String] = Task { "main.py" }
-  def sources: T[PathRef] = Task.Source(millSourcePath / "src")
+
+  /**
+   * The folders where the source files for this mill module live
+   *
+   * Python modules will be defined relative to these directories.
+   */
+  def sources: T[Seq[PathRef]] = Task.Sources { millSourcePath / "src" }
+
+  /**
+   * The script to run. This file may not exist if this module is only a library.
+   */
+  def script: T[PathRef] = Task.Source { millSourcePath / "src" / "main.py" }
 
   def pythonDeps: T[Seq[String]] = Task { Seq.empty[String] }
 
@@ -13,6 +27,11 @@ trait PythonModule extends Module {
     pythonDeps() ++ upstreamDependencies
   }
 
+  def transitiveSources: T[Seq[PathRef]] = Task {
+    val upstreamSources = Task.traverse(moduleDeps)(_.transitiveSources)().flatten
+    sources() ++ upstreamSources
+  }
+
   def pythonExe: T[PathRef] = Task {
     os.call(("python3", "-m", "venv", Task.dest / "venv"))
     val python = Task.dest / "venv" / "bin" / "python3"
@@ -25,55 +44,72 @@ trait PythonModule extends Module {
     Task.traverse(moduleDeps)(_.typeCheck)()
 
     os.call(
-      (pythonExe().path, "-m", "mypy", "--strict", sources().path),
+      (pythonExe().path, "-m", "mypy", "--strict", sources().map(_.path)),
       stdout = os.Inherit,
       cwd = T.workspace
     )
   }
 
-  def gatherScripts(upstream: Seq[(PathRef, PythonModule)]) = {
-    for ((sourcesFolder, mod) <- upstream) {
-      val destinationPath =
-        os.pwd / mod.millSourcePath.subRelativeTo(mill.api.WorkspaceRoot.workspaceRoot)
-      os.copy.over(sourcesFolder.path / os.up, destinationPath)
-    }
-  }
-
   def run(args: mill.define.Args) = Task.Command {
-    gatherScripts(Task.traverse(moduleDeps)(_.sources)().zip(moduleDeps))
-
     os.call(
-      (pythonExe().path, sources().path / mainFileName(), args.value),
-      env = Map("PYTHONPATH" -> Task.dest.toString),
-      stdout = os.Inherit
+      (pythonExe().path, script().path, args.value),
+      env = Map(
+        "PYTHONPATH" -> transitiveSources().map(_.path).mkString(":"),
+        "PYTHONPYCACHEPREFIX" -> (T.dest / "cache").toString
+      ),
+      stdout = os.Inherit,
+      cwd = T.dest
     )
   }
 
+  override def defaultCommandName(): String = "run"
+
+  /**
+   * Opens up a Python console with your module and all dependencies present,
+   * for you to test and operate your code interactively.
+   */
+  def console(): Command[Unit] = Task.Command(exclusive = true) {
+    if (!Util.isInteractive()) {
+      Result.Failure("console needs to be run with the -i/--interactive flag")
+    } else {
+      Jvm.runSubprocess(
+        Seq(pythonExe().path.toString),
+        envArgs = Map(
+          "PYTHONPATH" -> transitiveSources().map(_.path).mkString(":").toString,
+          "PYTHONPYCACHEPREFIX" -> (T.dest / "cache").toString
+        ),
+        workingDir = Task.dest
+      )
+      Result.Success(())
+    }
+  }
+
   /** Bundles the project into a single PEX executable(bundle.pex). */
   def bundle = Task {
-    gatherScripts(Task.traverse(moduleDeps)(_.sources)().zip(moduleDeps))
-
     val pexFile = Task.dest / "bundle.pex"
     os.call(
       (
+        // format: off
         pythonExe().path,
-        "-m",
-        "pex",
+        "-m", "pex",
         transitivePythonDeps(),
-        "-D",
-        Task.dest,
-        "-c",
-        sources().path / mainFileName(),
-        "-o",
-        pexFile,
-        "--scie",
-        "eager"
+        transitiveSources().flatMap(pr =>
+          Seq("-D", pr.path.toString)
+        ),
+        "--exe", script().path,
+        "-o", pexFile,
+        "--scie", "eager",
+        // format: on
       ),
-      env = Map("PYTHONPATH" -> Task.dest.toString),
-      stdout = os.Inherit
+      stdout = os.Inherit,
+      cwd = T.dest
     )
 
     PathRef(pexFile)
   }
 
+  trait PythonTests extends PythonModule {
+    override def moduleDeps: Seq[PythonModule] = Seq(outer)
+  }
+
 }

pythonlib/src/mill/pythonlib/TestModule.scala

@@ -0,0 +1,56 @@
+package mill.pythonlib
+
+import mill.Task
+import mill.Command
+import mill.TaskModule
+
+trait TestModule extends TaskModule {
+
+  // TODO: make this return something more interesting
+  def test(args: String*): Command[Unit]
+
+  override def defaultCommandName() = "test"
+}
+
+object TestModule {
+
+  trait Unittest extends PythonModule with TestModule {
+    def test(args: String*): Command[Unit] = Task.Command {
+      val testArgs = if (args.isEmpty) {
+        Seq("discover") ++ sources().flatMap(pr => Seq("-s", pr.path.toString))
+      } else {
+        args
+      }
+
+      os.call(
+        (pythonExe().path, "-m", "unittest", testArgs),
+        env = Map(
+          "PYTHONPATH" -> transitiveSources().map(_.path).mkString(":"),
+          "PYTHONPYCACHEPREFIX" -> (Task.dest / "cache").toString
+        ),
+        stdout = os.Inherit,
+        cwd = Task.dest
+      )
+      ()
+    }
+  }
+
+  trait Pytest extends PythonModule with TestModule {
+
+    def pythonDeps = Seq("pytest")
+
+    def test(args: String*): Command[Unit] = Task.Command {
+      os.call(
+        (pythonExe().path, "-m", "pytest", sources().map(_.path), args),
+        env = Map(
+          "PYTHONPATH" -> transitiveSources().map(_.path).mkString(":"),
+          "PYTHONPYCACHEPREFIX" -> (Task.dest / "cache").toString
+        ),
+        stdout = os.Inherit,
+        cwd = Task.dest
+      )
+      ()
+    }
+  }
+
+}