add environment variables support for PythonTask (#8)

Author: xvik

CHANGELOG.md

@@ -1,3 +1,7 @@
+* Add environment variables support (#8):
+    - Python object: `Python.environment` (single and map)
+    - Python task: `PythonTask.environment` (single and map)
+
 ### 2.0.0 (2020-03-13)
 * (breaking) Drop java 7 support
 * (breaking) Drop gradle 4 support

src/main/groovy/ru/vyarus/gradle/plugin/python/cmd/Python.groovy

@@ -44,6 +44,7 @@ class Python {
     private LogLevel logLevel = LogLevel.INFO
     private final List<String> pythonArgs = []
     private final List<String> extraArgs = []
+    private final Map<String, Object> envVars = [:]
 
     // run through cmd on win (when direct executable called)
     private final boolean withCmd
@@ -137,6 +138,37 @@ class Python {
         return this
     }
 
+    /**
+     * Set environment variable for executed python process.
+     * <p>
+     * Specified variables could be cleared with {@link #clearEnvironment()}.
+     *
+     * @param var environment variable
+     * @param value variable value
+     * @return cli instance for chained calls
+     */
+    Python environment(String var, Object value) {
+        envVars.put(var, value)
+        return this
+    }
+
+    /**
+     * Set environment variables for executed python process. This call does not replace previous variables map,
+     * it only adds all specified values into existing map (with possible overriding of some variable values).
+     * So if called multiple times, all specified maps would be aggregated.
+     * <p>
+     * Specified variables could be cleared with {@link #clearEnvironment()}.
+     *
+     * @param vars environment variables (may be null)
+     * @return cli instance for chained calls
+     */
+    Python environment(Map<String, Object> vars) {
+        if (vars) {
+            envVars.putAll(vars)
+        }
+        return this
+    }
+
     /**
      * Removes all registered python arguments.
      *
@@ -157,6 +189,11 @@ class Python {
         return this
     }
 
+    Python clearEnvironment() {
+        this.envVars.clear()
+        return this
+    }
+
     /**
      * Execute python command and return all output (without applying the prefix!).
      * Very handy for short scripts evaluation. In case of error, all output is logged (with prefix).
@@ -327,6 +364,9 @@ class Python {
             if (workDir) {
                 setWorkingDir(workDir)
             }
+            if (envVars) {
+                it.environment(envVars)
+            }
         }
         project.logger.info('Python execution time: {}',
                 DurationFormatter.format(System.currentTimeMillis() - start))

src/main/groovy/ru/vyarus/gradle/plugin/python/task/PythonTask.groovy

@@ -75,6 +75,13 @@ class PythonTask extends BasePythonTask {
     @Input
     @Optional
     List<String> extraArgs = []
+    /**
+     * Environment variables for executed python process (variables specified in gradle's
+     * {@link org.gradle.process.ExecSpec#environment(java.util.Map)} during python process execution).
+     */
+    @Input
+    @Optional
+    Map<String, Object> environment = [:]
     /**
      * Prefix each line of python output. By default it's '\t' to indicate command output.
      */
@@ -97,6 +104,7 @@ class PythonTask extends BasePythonTask {
                 .workDir(getWorkDir())
                 .pythonArgs(getPythonArgs())
                 .extraArgs(getExtraArgs())
+                .environment(environment)
 
         if (mod) {
             python.callModule(mod, cmd)
@@ -129,6 +137,30 @@ class PythonTask extends BasePythonTask {
         }
     }
 
+    /**
+     * Add environment variable for python process (will override previously set value).
+     *
+     * @param var variable name
+     * @param value variable value
+     */
+    @SuppressWarnings('ConfusingMethodName')
+    void environment(String var, Object value) {
+        getEnvironment().put(var, value)
+    }
+
+    /**
+     * Add environment variables for python process (will override already set values, but not replace context
+     * map completely). May be called multiple times: all variables would be aggregated.
+     *
+     * @param vars (may be null)
+     */
+    @SuppressWarnings('ConfusingMethodName')
+    void environment(Map<String, Object> vars) {
+        if (vars) {
+            getEnvironment().putAll(vars)
+        }
+    }
+
     private void initWorkDirIfRequired() {
         String dir = getWorkDir()
         if (dir && isCreateWorkDir()) {

src/test/groovy/ru/vyarus/gradle/plugin/python/cmd/PythonCliTest.groovy

@@ -125,5 +125,28 @@ class PythonCliTest extends AbstractTest {
         python.clearExtraArgs()
         then: 'set'
         python.extraArgs.empty
+
+        when: "set vars"
+        python.environment('foo', 1)
+        python.environment('bar', 2)
+        then: 'set'
+        python.envVars == ['foo': 1, 'bar' : 2]
+
+        when: 'clear vars'
+        python.clearEnvironment()
+        then: 'no vars'
+        python.envVars.size() == 0
+
+        when: 'mass variables set'
+        python.environment(['foo': 2, 'bar' : 3])
+        python.environment(['foo': 1, 'baz' : 4])
+        then: 'aggregated'
+        python.envVars == ['foo': 1, 'bar' : 3, 'baz': 4]
+
+        when: 'additional var'
+        python.environment('sample': 'sam')
+        then: 'aggregated'
+        python.envVars == ['foo': 1, 'bar' : 3, 'baz': 4, 'sample': 'sam']
+
     }
 }

src/test/groovy/ru/vyarus/gradle/plugin/python/task/PythonTaskKitTest.groovy

@@ -184,4 +184,52 @@ class PythonTaskKitTest extends AbstractKitTest {
         result.output.contains('\t sample')
     }
 
+
+    def "Check env vars"() {
+        setup:
+        build """
+            plugins {
+                id 'ru.vyarus.use-python'
+            }
+
+            task sample(type: PythonTask) {
+                command = "-c \\"import os;print('variables: '+os.getenv('some', 'null')+' '+os.getenv('foo', 'null'))\\""
+                environment 'some', 1
+                environment(['foo': 'bar'])
+            }
+        """
+
+        when: "run task"
+        debug()
+        BuildResult result = run('sample')
+
+        then: "executed"
+        result.task(':sample').outcome == TaskOutcome.SUCCESS
+        result.output.contains('variables: 1 bar')
+    }
+
+
+    def "Check python see system variables"() {
+        setup:
+        build """
+            plugins {
+                id 'ru.vyarus.use-python'
+            }
+            
+            assert System.getenv('some') == 'foo'
+            
+            task sample(type: PythonTask) {
+                command = "-c \\"import os;print('variables: '+os.getenv('some', 'null'))\\""
+            }
+        """
+
+        when: "run task"
+        BuildResult result = gradle('sample')
+                .withEnvironment(['some': 'foo', 'PATH': System.getenv('PATH')])
+                .build()
+
+        then: "executed"
+        result.task(':sample').outcome == TaskOutcome.SUCCESS
+        result.output.contains('variables: foo')
+    }
 }