Add prompt parameter to python::pyvenv

Add a prompt parameter to python::pyvenv that uses the --prompt command-line argument available in the venv module in Python 3.6 and later, to set the prompt shown when the virtualenv is active.
voxpupuli · Sep 1, 2022 · 343e97d · 343e97d
1 parent c47ce1e
commit 343e97d
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 4 deletions.
diff --git a/REFERENCE.md b/REFERENCE.md
@@ -903,6 +903,7 @@ The following parameters are available in the `python::pyvenv` defined type:
 * [`mode`](#mode)
 * [`path`](#path)
 * [`environment`](#environment)
+* [`prompt`](#prompt)
 * [`pip_version`](#pip_version)
 
 ##### <a name="ensure"></a>`ensure`
@@ -977,6 +978,14 @@ Optionally specify environment variables for pyvenv
 
 Default value: `[]`
 
+##### <a name="prompt"></a>`prompt`
+
+Data type: `Variant[Boolean,String[1]]`
+
+Optionally specify the virtualenv prompt (python >= 3.6)
+
+Default value: ``false``
+
 ##### <a name="pip_version"></a>`pip_version`
 
 Data type: `Python::Venv::PipVersion`

diff --git a/manifests/pyvenv.pp b/manifests/pyvenv.pp
@@ -10,6 +10,7 @@
 # @param mode  Optionally specify directory mode
 # @param path Specifies the PATH variable.
 # @param environment Optionally specify environment variables for pyvenv
+# @param prompt Optionally specify the virtualenv prompt (python >= 3.6)
 #
 # @example
 #   python::pyvenv { '/var/www/project1' :
@@ -31,6 +32,7 @@
   Stdlib::Filemode            $mode        = '0755',
   Array[Stdlib::Absolutepath] $path        = ['/bin', '/usr/bin', '/usr/sbin', '/usr/local/bin',],
   Array                       $environment = [],
+  Variant[Boolean,String[1]]  $prompt      = false,
   Python::Venv::PipVersion    $pip_version = 'latest',
 ) {
   include python
@@ -62,6 +64,12 @@
       $system_pkgs_flag = ''
     }
 
+    if versioncmp($normalized_python_version, '3.6') >=0 and $prompt {
+      $prompt_arg = "--prompt ${shell_escape($prompt)}"
+    } else {
+      $prompt_arg = ''
+    }
+
     file { $venv_dir:
       ensure  => directory,
       owner   => $owner,
@@ -78,7 +86,7 @@
     }
 
     exec { "python_virtualenv_${venv_dir}":
-      command     => "${virtualenv_cmd} --clear ${system_pkgs_flag} ${venv_dir} && ${pip_cmd} --log ${venv_dir}/pip.log install ${pip_upgrade} && ${pip_cmd} --log ${venv_dir}/pip.log install --upgrade setuptools",
+      command     => "${virtualenv_cmd} --clear ${system_pkgs_flag} ${prompt_arg} ${venv_dir} && ${pip_cmd} --log ${venv_dir}/pip.log install ${pip_upgrade} && ${pip_cmd} --log ${venv_dir}/pip.log install --upgrade setuptools",
       user        => $owner,
       creates     => "${venv_dir}/bin/activate",
       path        => $_path,

diff --git a/spec/classes/python_spec.rb b/spec/classes/python_spec.rb
@@ -146,7 +146,7 @@
               it {
                 expect(subject).to contain_exec('python_virtualenv_/opt/env1').
                   with(
-                    command: 'python3.8 -m venv --clear  /opt/env1 && /opt/env1/bin/pip --log /opt/env1/pip.log install --upgrade pip && /opt/env1/bin/pip --log /opt/env1/pip.log install --upgrade setuptools',
+                    command: 'python3.8 -m venv --clear   /opt/env1 && /opt/env1/bin/pip --log /opt/env1/pip.log install --upgrade pip && /opt/env1/bin/pip --log /opt/env1/pip.log install --upgrade setuptools',
                     user: 'root',
                     creates: '/opt/env1/bin/activate',
                     path: [
@@ -166,7 +166,7 @@
               it {
                 expect(subject).to contain_exec('python_virtualenv_/opt/env2').
                   with(
-                    command: 'python3.8 -m venv --clear  /opt/env2 && /opt/env2/bin/pip --log /opt/env2/pip.log install --upgrade \'pip <= 20.3.4\' && /opt/env2/bin/pip --log /opt/env2/pip.log install --upgrade setuptools',
+                    command: 'python3.8 -m venv --clear   /opt/env2 && /opt/env2/bin/pip --log /opt/env2/pip.log install --upgrade \'pip <= 20.3.4\' && /opt/env2/bin/pip --log /opt/env2/pip.log install --upgrade setuptools',
                     user: 'root',
                     creates: '/opt/env2/bin/activate',
                     path: [

diff --git a/spec/defines/pyvenv_spec.rb b/spec/defines/pyvenv_spec.rb
@@ -19,7 +19,7 @@
 
       context 'with default parameters' do
         it { is_expected.to contain_file('/opt/env').that_requires('Class[python::install]') }
-        it { is_expected.to contain_exec('python_virtualenv_/opt/env').with_command('pyvenv-3.5 --clear  /opt/env && /opt/env/bin/pip --log /opt/env/pip.log install --upgrade pip && /opt/env/bin/pip --log /opt/env/pip.log install --upgrade setuptools') }
+        it { is_expected.to contain_exec('python_virtualenv_/opt/env').with_command('pyvenv-3.5 --clear   /opt/env && /opt/env/bin/pip --log /opt/env/pip.log install --upgrade pip && /opt/env/bin/pip --log /opt/env/pip.log install --upgrade setuptools') }
       end
 
       describe 'when ensure' do
@@ -36,5 +36,30 @@
         end
       end
     end
+
+    context "prompt on #{os} with python 3.6" do
+      let :facts do
+        # python 3.6 is required for venv and prompt
+        facts.merge(
+          python3_version: '3.6.1'
+        )
+      end
+      let :title do
+        '/opt/env'
+      end
+
+      context 'with prompt' do
+        let :params do
+          {
+            prompt: 'custom prompt',
+          }
+        end
+
+        it {
+          is_expected.to contain_file('/opt/env').that_requires('Class[python::install]')
+          is_expected.to contain_exec('python_virtualenv_/opt/env').with_command('python3.6 -m venv --clear  --prompt custom\\ prompt /opt/env && /opt/env/bin/pip --log /opt/env/pip.log install --upgrade pip && /opt/env/bin/pip --log /opt/env/pip.log install --upgrade setuptools')
+        }
+      end
+    end
   end
 end