Merge pull request #1616 from alexarchambault/current-dir-python-path

alexarchambault · web-flow · commit b66141b560a1 · 2023-01-17T11:43:09.000+01:00
Put workspace and / or current directory in PYTHONPATH in the run and repl commands
diff --git a/build.sc b/build.sc
@@ -194,6 +194,13 @@ object dummy extends Module {
       Deps.pythonInterface
     )
   }
+  object scalaPy extends ScalaModule with Bloop.Module {
+    def skipBloop    = true
+    def scalaVersion = Scala.defaultInternal
+    def ivyDeps = Agg(
+      Deps.scalaPy
+    )
+  }
 }
 
 trait BuildMacros extends ScalaCliSbtModule
diff --git a/modules/cli/src/main/scala/scala/cli/commands/repl/Repl.scala b/modules/cli/src/main/scala/scala/cli/commands/repl/Repl.scala
@@ -14,7 +14,11 @@ import scala.build.internal.{Constants, Runner}
 import scala.build.options.{BuildOptions, JavaOpt, MaybeScalaVersion, Scope}
 import scala.cli.CurrentParams
 import scala.cli.commands.publish.ConfigUtil.*
-import scala.cli.commands.run.Run.{maybePrintSimpleScalacOutput, orPythonDetectionError}
+import scala.cli.commands.run.Run.{
+  maybePrintSimpleScalacOutput,
+  orPythonDetectionError,
+  pythonPathEnv
+}
 import scala.cli.commands.run.RunMode
 import scala.cli.commands.shared.SharedOptions
 import scala.cli.commands.{ScalaCommand, WatchUtil}
@@ -264,20 +268,25 @@ object Repl extends ScalaCommand[ReplOptions] {
       }
       .params
 
-    val scalapyJavaOpts =
+    val (scalapyJavaOpts, scalapyExtraEnv) =
       if (setupPython) {
         val props = value {
           val python       = Python()
           val propsOrError = python.scalapyProperties
           logger.debug(s"Python Java properties: $propsOrError")
           propsOrError.orPythonDetectionError
         }
-        props.toVector.sorted.map {
+        val props0 = props.toVector.sorted.map {
           case (k, v) => s"-D$k=$v"
         }
+        // Putting current dir in PYTHONPATH, see
+        // https://github.com/VirtusLab/scala-cli/pull/1616#issuecomment-1333283174
+        // for context.
+        val dirs = buildOpt.map(_.inputs.workspace).toSeq ++ Seq(os.pwd)
+        (props0, pythonPathEnv(dirs: _*))
       }
       else
-        Nil
+        (Nil, Map.empty[String, String])
 
     def additionalArgs = {
       val pythonArgs =
@@ -365,7 +374,7 @@ object Repl extends ScalaCommand[ReplOptions] {
           maybeAdaptForWindows(replArgs),
           logger,
           allowExecve = allowExit,
-          extraEnv = extraEnv
+          extraEnv = scalapyExtraEnv ++ extraEnv
         ).waitFor()
         if (retCode != 0)
           value(Left(new ReplError(retCode)))
diff --git a/modules/cli/src/main/scala/scala/cli/commands/run/Run.scala b/modules/cli/src/main/scala/scala/cli/commands/run/Run.scala
@@ -4,6 +4,7 @@ import ai.kien.python.Python
 import caseapp.*
 
 import java.io.File
+import java.util.Locale
 import java.util.concurrent.CompletableFuture
 
 import scala.build.EitherCps.{either, value}
@@ -338,6 +339,27 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
     value(res)
   }
 
+  def pythonPathEnv(dirs: os.Path*): Map[String, String] = {
+    val onlySafePaths = sys.env.exists {
+      case (k, v) =>
+        k.toLowerCase(Locale.ROOT) == "pythonsafepath" && v.nonEmpty
+    }
+    // Don't add unsafe directories to PYTHONPATH if PYTHONSAFEPATH is set,
+    // see https://docs.python.org/3/using/cmdline.html#envvar-PYTHONSAFEPATH
+    // and https://github.com/VirtusLab/scala-cli/pull/1616#issuecomment-1336017760
+    // for more details.
+    if (onlySafePaths) Map.empty[String, String]
+    else {
+      val (pythonPathEnvVarName, currentPythonPath) = sys.env
+        .find(_._1.toLowerCase(Locale.ROOT) == "pythonpath")
+        .getOrElse(("PYTHONPATH", ""))
+      val updatedPythonPath = (currentPythonPath +: dirs.map(_.toString))
+        .filter(_.nonEmpty)
+        .mkString(File.pathSeparator)
+      Map(pythonPathEnvVarName -> updatedPythonPath)
+    }
+  }
+
   private def runOnce(
     build: Build.Successful,
     mainClass: String,
@@ -397,9 +419,9 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
         value(res)
       case Platform.Native =>
         val setupPython = build.options.notForBloopOptions.doSetupPython.getOrElse(false)
-        val (pythonExecutable, pythonLibraryPaths) =
-          if (setupPython)
-            value {
+        val (pythonExecutable, pythonLibraryPaths, pythonExtraEnv) =
+          if (setupPython) {
+            val (exec, libPaths) = value {
               val python = Python()
               val pythonPropertiesOrError = for {
                 paths      <- python.nativeLibraryPaths
@@ -408,8 +430,13 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
               logger.debug(s"Python executable and native library paths: $pythonPropertiesOrError")
               pythonPropertiesOrError.orPythonDetectionError
             }
+            // Putting the workspace in PYTHONPATH, see
+            // https://github.com/VirtusLab/scala-cli/pull/1616#issuecomment-1333283174
+            // for context.
+            (exec, libPaths, pythonPathEnv(build.inputs.workspace))
+          }
           else
-            (None, Nil)
+            (None, Nil, Map())
         // seems conda doesn't add the lib directory to LD_LIBRARY_PATH (see conda/conda#308),
         // which prevents apps from finding libpython for example, so we update it manually here
         val libraryPathsEnv =
@@ -434,7 +461,7 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
           }
         val programNameEnv =
           pythonExecutable.fold(Map.empty)(py => Map("SCALAPY_PYTHON_PROGRAMNAME" -> py))
-        val extraEnv = libraryPathsEnv ++ programNameEnv
+        val extraEnv = libraryPathsEnv ++ programNameEnv ++ pythonExtraEnv
         val maybeResult = withNativeLauncher(
           build,
           mainClass,
@@ -463,20 +490,24 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
           case RunMode.Default =>
             val baseJavaProps = build.options.javaOptions.javaOpts.toSeq.map(_.value.value)
             val setupPython   = build.options.notForBloopOptions.doSetupPython.getOrElse(false)
-            val pythonJavaProps =
+            val (pythonJavaProps, pythonExtraEnv) =
               if (setupPython) {
                 val scalapyProps = value {
                   val python       = Python()
                   val propsOrError = python.scalapyProperties
                   logger.debug(s"Python Java properties: $propsOrError")
                   propsOrError.orPythonDetectionError
                 }
-                scalapyProps.toVector.sorted.map {
+                val props = scalapyProps.toVector.sorted.map {
                   case (k, v) => s"-D$k=$v"
                 }
+                // Putting the workspace in PYTHONPATH, see
+                // https://github.com/VirtusLab/scala-cli/pull/1616#issuecomment-1333283174
+                // for context.
+                (props, pythonPathEnv(build.inputs.workspace))
               }
               else
-                Nil
+                (Nil, Map.empty[String, String])
             val allJavaOpts = pythonJavaProps ++ baseJavaProps
             if (showCommand) {
               val command = Runner.jvmCommand(
@@ -485,6 +516,7 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
                 build.fullClassPath,
                 mainClass,
                 args,
+                extraEnv = pythonExtraEnv,
                 useManifest = build.options.notForBloopOptions.runWithManifest,
                 scratchDirOpt = scratchDirOpt
               )
@@ -499,6 +531,7 @@ object Run extends ScalaCommand[RunOptions] with BuildCommandHelpers {
                 args,
                 logger,
                 allowExecve = allowExecve,
+                extraEnv = pythonExtraEnv,
                 useManifest = build.options.notForBloopOptions.runWithManifest,
                 scratchDirOpt = scratchDirOpt
               )
diff --git a/modules/integration/src/test/scala/scala/cli/integration/ReplTestDefinitions.scala b/modules/integration/src/test/scala/scala/cli/integration/ReplTestDefinitions.scala
@@ -46,12 +46,18 @@ abstract class ReplTestDefinitions(val scalaVersionOpt: Option[String])
   }
 
   test("ammonite scalapy") {
-    TestInputs.empty.fromRoot { root =>
+    val inputs = TestInputs(
+      os.rel / "foo" / "something.py" ->
+        """messageStart = 'Hello from'
+          |messageEnd = 'ScalaPy'
+          |""".stripMargin
+    )
+    inputs.fromRoot { root =>
       val ammArgs = Seq(
         "-c",
         """println("Hello" + " from Scala " + scala.util.Properties.versionNumberString)
-          |// py.Dynamic.global.print("Hello from", "ScalaPy") // doesn't work
-          |println(py"'Hello from '" + py"'ScalaPy'")
+          |val sth = py.module("foo.something")
+          |py.Dynamic.global.applyDynamicNamed("print")("" -> sth.messageStart, "" -> sth.messageEnd, "flush" -> py.Any.from(true))
           |""".stripMargin
       )
         .map {
@@ -61,6 +67,24 @@ abstract class ReplTestDefinitions(val scalaVersionOpt: Option[String])
             identity
         }
         .flatMap(arg => Seq("--ammonite-arg", arg))
+
+      val errorRes = os.proc(
+        TestUtil.cli,
+        "repl",
+        extraOptions,
+        "--ammonite",
+        "--python",
+        ammArgs
+      ).call(
+        cwd = root,
+        env = Map("PYTHONSAFEPATH" -> "foo"),
+        mergeErrIntoOut = true,
+        check = false
+      )
+      expect(errorRes.exitCode != 0)
+      val errorOutput = errorRes.out.text()
+      expect(errorOutput.contains("No module named 'foo'"))
+
       val res = os.proc(
         TestUtil.cli,
         "repl",
diff --git a/modules/integration/src/test/scala/scala/cli/integration/RunScalaPyTestDefinitions.scala b/modules/integration/src/test/scala/scala/cli/integration/RunScalaPyTestDefinitions.scala
@@ -78,4 +78,61 @@ trait RunScalaPyTestDefinitions { _: RunTestDefinitions =>
     test("scalapy native") {
       scalapyNativeTest()
     }
+
+  def pythonAndScalaSourcesTest(native: Boolean): Unit = {
+    val tq = "\"\"\""
+    val inputs = TestInputs(
+      os.rel / "src" / "helpers.py" ->
+        s"""class Helper:
+           |    ${tq}Helper class$tq
+           |
+           |    def message(self):
+           |        return 'Hello from Python'
+           |""".stripMargin,
+      os.rel / "src" / "Hello.scala" ->
+        s"""//> using python
+           |$maybeScalapyPrefix
+           |object Hello {
+           |  def main(args: Array[String]): Unit =
+           |    py.local {
+           |      val helpers = py.module("helpers")
+           |      println(helpers.Helper().message())
+           |    }
+           |}
+           |""".stripMargin
+    )
+    val nativeOpt = if (native) Seq("--native") else Nil
+    inputs.fromRoot { root =>
+
+      // Script dir shouldn't be added to PYTHONPATH if PYTHONSAFEPATH is non-empty
+      val errorRes = os.proc(TestUtil.cli, "run", extraOptions, nativeOpt, "src")
+        .call(
+          cwd = root,
+          env = Map("PYTHONSAFEPATH" -> "foo"),
+          mergeErrIntoOut = true,
+          check = false
+        )
+      expect(errorRes.exitCode != 0)
+      val errorOutput = errorRes.out.text()
+      expect(errorOutput.contains("No module named 'helpers'"))
+
+      val res = os.proc(TestUtil.cli, "run", extraOptions, nativeOpt, "src")
+        .call(cwd = root)
+      val output = res.out.trim()
+      if (native)
+        expect(output.linesIterator.toVector.endsWith(Seq("Hello from Python")))
+      else
+        expect(output == "Hello from Python")
+    }
+  }
+
+  test("Python and Scala sources") {
+    pythonAndScalaSourcesTest(native = false)
+  }
+  // disabled on Windows for now, for context, see
+  // https://github.com/VirtusLab/scala-cli/pull/1270#issuecomment-1237904394
+  if (!Properties.isWin)
+    test("Python and Scala sources (native)") {
+      pythonAndScalaSourcesTest(native = true)
+    }
 }
diff --git a/project/deps.sc b/project/deps.sc
@@ -160,7 +160,7 @@ object Deps {
   def scalaPackagerCli = ivy"org.virtuslab:scala-packager-cli_2.13:${Versions.scalaPackager}"
   // Force using of 2.13 - is there a better way?
   def scalaparse               = ivy"com.lihaoyi:scalaparse_2.13:2.3.3"
-  def scalaPy                  = ivy"me.shadaj::scalapy-core::0.5.2+5-83f1eb68"
+  def scalaPy                  = ivy"dev.scalapy::scalapy-core::0.5.3"
   def scalaReflect(sv: String) = ivy"org.scala-lang:scala-reflect:$sv"
   def semanticDbJavac          = ivy"com.sourcegraph:semanticdb-javac:0.7.4"
   def semanticDbScalac         = ivy"org.scalameta:::semanticdb-scalac:${Versions.scalaMeta}"
diff --git a/website/docs/reference/cli-options.md b/website/docs/reference/cli-options.md
@@ -1039,7 +1039,7 @@ Enable Python support via ScalaPy
 
 Aliases: `--scalapy-version`
 
-[experimental] Set ScalaPy version (0.5.2+5-83f1eb68 by default)
+[experimental] Set ScalaPy version (0.5.3 by default)
 
 ## Repl options
 

Original file line number	Diff line number	Diff line change
`@@ -194,6 +194,13 @@ object dummy extends Module {`
`194`	`194`	`Deps.pythonInterface`
`195`	`195`	`)`
`196`	`196`	`}`
	`197`	`+ object scalaPy extends ScalaModule with Bloop.Module {`
	`198`	`+ def skipBloop = true`
	`199`	`+ def scalaVersion = Scala.defaultInternal`
	`200`	`+ def ivyDeps = Agg(`
	`201`	`+ Deps.scalaPy`
	`202`	`+ )`
	`203`	`+ }`
`197`	`204`	`}`
`198`	`205`
`199`	`206`	`trait BuildMacros extends ScalaCliSbtModule`