Some improvements to pythonlib

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()

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 {
   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,52 +44,65 @@ 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)