diff --git a/powerline/bindings/zsh/__init__.py b/powerline/bindings/zsh/__init__.py index 1037bba2..44f5c697 100644 --- a/powerline/bindings/zsh/__init__.py +++ b/powerline/bindings/zsh/__init__.py @@ -101,10 +101,7 @@ class Environment(object): return False -if hasattr(getattr(zsh, 'environ', None), '__contains__'): - environ = zsh.environ -else: - environ = Environment() +environ = Environment() if hasattr(zsh, 'expand') and zsh.expand('${:-}') == '': diff --git a/powerline/segments/common/env.py b/powerline/segments/common/env.py index 61f516e2..bbfe3e25 100644 --- a/powerline/segments/common/env.py +++ b/powerline/segments/common/env.py @@ -19,20 +19,24 @@ def environment(pl, segment_info, variable=None): @requires_segment_info -def virtualenv(pl, segment_info, ignore_venv=False, ignore_conda=False): +def virtualenv(pl, segment_info, ignore_venv=False, ignore_conda=False, ignored_names=("venv", ".venv")): '''Return the name of the current Python or conda virtualenv. - + :param list ignored_names: + Names of venvs to ignore. Will then get the name of the venv by ascending to the parent directory :param bool ignore_venv: Whether to ignore virtual environments. Default is False. :param bool ignore_conda: Whether to ignore conda environments. Default is False. ''' - return ( - (not ignore_venv and - os.path.basename(segment_info['environ'].get('VIRTUAL_ENV', ''))) or - (not ignore_conda and - segment_info['environ'].get('CONDA_DEFAULT_ENV', '')) or - None) + if not ignore_venv: + for candidate in reversed(segment_info['environ'].get('VIRTUAL_ENV', '').split("/")): + if candidate and candidate not in ignored_names: + return candidate + if not ignore_conda: + for candidate in reversed(segment_info['environ'].get('CONDA_DEFAULT_ENV', '').split("/")): + if candidate and candidate not in ignored_names: + return candidate + return None @requires_segment_info diff --git a/tests/test_python/test_segments.py b/tests/test_python/test_segments.py index 7564bff0..6c441143 100644 --- a/tests/test_python/test_segments.py +++ b/tests/test_python/test_segments.py @@ -687,6 +687,10 @@ class TestEnv(TestCommon): self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignore_conda=True), 'ghi') self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignore_venv=True), None) self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignore_venv=True, ignore_conda=True), None) + self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignored_names=["aaa"]), "ghi") + self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignored_names=["ghi"]), "def") + self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignored_names=["def", "ghi"]), "abc") + self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignored_names=["abc", "def", "ghi"]), None) segment_info['environ'].pop('VIRTUAL_ENV') self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info), None) @@ -696,6 +700,7 @@ class TestEnv(TestCommon): with replace_env('CONDA_DEFAULT_ENV', 'foo') as segment_info: self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info), 'foo') + self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignored_names=["foo"]), None) self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignore_conda=True), None) self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignore_venv=True), 'foo') self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignore_venv=True, ignore_conda=True), None) @@ -718,6 +723,9 @@ class TestEnv(TestCommon): self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignore_venv=True), None) self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info, ignore_venv=True, ignore_conda=True), None) + with replace_env('VIRTUAL_ENV', '/abc/def/venv') as segment_info: + self.assertEqual(self.module.virtualenv(pl=pl, segment_info=segment_info), 'def') + def test_environment(self): pl = Pl() variable = 'FOO'