Extended environment variable support for session and server.

Author: tasdomas

tmuxp/common.py

@@ -0,0 +1,103 @@
+
+class EnvironmentMixin(object):
+
+    """Mixin class for managing session and server level environment
+    variables in tmux.
+
+    """
+
+    _add_option = None
+
+    def __init__(self, add_option=None):
+        self._add_option = add_option
+    
+    def set_environment(self, name, value):
+        """Set environment ``$ tmux set-environment <name> <value>``.
+
+        :param name: the environment variable name. such as 'PATH'.
+        :type option: string
+        :param value: environment value.
+        :type value: string
+
+        """
+
+        args = ['set-environment']
+        if self._add_option:
+            args += [self._add_option]
+
+        args += [name, value]
+            
+        proc = self.cmd(*args)
+
+        if proc.stderr:
+            if isinstance(proc.stderr, list) and len(proc.stderr) == int(1):
+                proc.stderr = proc.stderr[0]
+            raise ValueError('tmux set-environment stderr: %s' % proc.stderr)
+
+    def unset_environment(self, name):
+        """Unset environment variable ``$ tmux set-environment -u <name>``.
+
+        :param name: the environment variable name. such as 'PATH'.
+        :type option: string
+        """
+
+        args = ['set-environment']
+        if self._add_option:
+            args += [self._add_option]
+        args += ['-u', name]
+            
+        proc = self.cmd(*args)
+
+        if proc.stderr:
+            if isinstance(proc.stderr, list) and len(proc.stderr) == int(1):
+                proc.stderr = proc.stderr[0]
+            raise ValueError('tmux set-environment stderr: %s' % proc.stderr)
+
+    def remove_environment(self, name):
+        """Remove environment variable ``$ tmux set-environment -r <name>``.
+
+        :param name: the environment variable name. such as 'PATH'.
+        :type option: string
+        """
+
+        args = ['set-environment']
+        if self._add_option:
+            args += [self._add_option]
+        args += ['-r', name]
+            
+        proc = self.cmd(*args)
+
+        if proc.stderr:
+            if isinstance(proc.stderr, list) and len(proc.stderr) == int(1):
+                proc.stderr = proc.stderr[0]
+            raise ValueError('tmux set-environment stderr: %s' % proc.stderr)
+
+    def show_environment(self, name=None):
+        """Show environment ``$tmux show-environment -t [session] <name>``.
+
+        Return dict of environment variables for the session or the value of a
+        specific variable if the name is specified.
+
+        :param name: the environment variable name. such as 'PATH'.
+        :type option: string
+        """
+        tmux_args = ['show-environment']
+        if self._add_option:
+            tmux_args += [self._add_option]
+        if name:
+            tmux_args += [name]
+        vars = self.cmd(*tmux_args).stdout
+        vars = [tuple(item.split('=')) for item in vars]
+        vars_dict = {}
+        for t in vars:
+            if len(t) == 2:
+                vars_dict[t[0]] = t[1]
+            elif len(t) == 1:
+                vars_dict[t[0]] = True
+            else:
+                raise ValueError('unexpected variable %s', t)
+
+        if name:
+            return vars_dict.get(name)
+
+        return vars_dict

tmuxp/server.py

@@ -14,11 +14,12 @@
 from . import formats, exc
 from .session import Session
 from .util import tmux_cmd, TmuxRelationalObject
+from .common import EnvironmentMixin
 
 logger = logging.getLogger(__name__)
 
 
-class Server(TmuxRelationalObject):
+class Server(TmuxRelationalObject, EnvironmentMixin):
 
     """The :term:`tmux(1)` server.
 
@@ -54,6 +55,7 @@ def __init__(
         colors=None,
         **kwargs
     ):
+        EnvironmentMixin.__init__(self, '-g')
         self._windows = []
         self._panes = []
 
@@ -383,7 +385,7 @@ def attach_session(self, target_session=None):
 
         if proc.stderr:
             raise exc.TmuxpException(proc.stderr)
-
+        
     def new_session(self,
                     session_name=None,
                     kill_session=False,

tmuxp/session.py

@@ -13,11 +13,12 @@
 
 from . import util, formats, exc
 from .window import Window
+from .common import EnvironmentMixin
 
 logger = logging.getLogger(__name__)
 
 
-class Session(util.TmuxMappingObject, util.TmuxRelationalObject):
+class Session(util.TmuxMappingObject, util.TmuxRelationalObject, EnvironmentMixin):
 
     """:term:`tmux(1)` session.
 
@@ -28,7 +29,7 @@ class Session(util.TmuxMappingObject, util.TmuxRelationalObject):
     childIdAttribute = 'window_id'
 
     def __init__(self, server=None, **kwargs):
-
+        EnvironmentMixin.__init__(self)
         self.server = server
 
         if not 'session_id' in kwargs:
@@ -330,27 +331,6 @@ def set_option(self, option, value):
                 proc.stderr = proc.stderr[0]
             raise ValueError('tmux set-option stderr: %s' % proc.stderr)
 
-    def set_environment(self, name, value):
-        """Set environment ``$ tmux set-environment <name> <value>``.
-
-        todo: needs tests
-
-        :param name: the environment variable name. such as 'PATH'.
-        :type option: string
-        :param value: environment value.
-        :type value: string
-
-        """
-
-        proc = self.cmd(
-            'set-environment', name, value
-        )
-
-        if proc.stderr:
-            if isinstance(proc.stderr, list) and len(proc.stderr) == int(1):
-                proc.stderr = proc.stderr[0]
-            raise ValueError('tmux set-environment stderr: %s' % proc.stderr)
-
     def show_options(self, option=None, g=False):
         """Return a dict of options for the window.
 

tmuxp/testsuite/server.py

@@ -67,6 +67,29 @@ def test_88_colors(self):
         self.assertIn('-8', proc.cmd)
         self.assertNotIn('-2', proc.cmd)
 
+class Environment(TmuxTestCase):
+
+    def test_show_environment(self):
+        """Server.show_environment() returns dict."""
+
+        vars = self.server.show_environment()
+        self.assertIsInstance(vars, dict)
+
+    def test_set_show_environment_single(self):
+        """Set environment then Server.show_environment(key)."""
+
+        self.server.set_environment('FOO', 'BAR')
+        self.assertEqual('BAR', self.server.show_environment('FOO'))
+
+        self.server.set_environment('FOO', 'DAR')
+        self.assertEqual('DAR', self.server.show_environment('FOO'))
+
+        self.assertEqual('DAR', self.server.show_environment()['FOO'])
+
+    def test_show_environment_not_set(self):
+        """Unset environment variable returns None."""
+        self.assertEqual(None, self.server.show_environment('BAR'))
+
 
 def suite():
     suite = unittest.TestSuite()

tmuxp/testsuite/session.py

@@ -140,9 +140,50 @@ def test_set_option_bad(self):
             self.session.set_option('afewewfew', 43)
 
 
+class Environment(TmuxTestCase):
+
+    def test_show_environment(self):
+        """Session.show_environment() returns dict."""
+
+        vars = self.session.show_environment()
+        self.assertIsInstance(vars, dict)
+
+    def test_set_show_environment_single(self):
+        """Set environment then Session.show_environment(key)."""
+
+        self.session.set_environment('FOO', 'BAR')
+        self.assertEqual('BAR', self.session.show_environment('FOO'))
+
+        self.session.set_environment('FOO', 'DAR')
+        self.assertEqual('DAR', self.session.show_environment('FOO'))
+
+        self.assertEqual('DAR', self.session.show_environment()['FOO'])
+
+    def test_show_environment_not_set(self):
+        """Not set environment variable returns None."""
+        self.assertEqual(None, self.session.show_environment('BAR'))
+
+    def test_remove_environment(self):
+        """Remove environment variable."""
+        self.assertEqual(None, self.session.show_environment('BAM'))
+        self.session.set_environment('BAM', 'OK')
+        self.assertEqual('OK', self.session.show_environment('BAM'))
+        self.session.remove_environment('BAM')
+        self.assertEqual(None, self.session.show_environment('BAM'))
+
+    def test_unset_environment(self):
+        """Unset environment variable."""
+        self.assertEqual(None, self.session.show_environment('BAM'))
+        self.session.set_environment('BAM', 'OK')
+        self.assertEqual('OK', self.session.show_environment('BAM'))
+        self.session.unset_environment('BAM')
+        self.assertEqual(None, self.session.show_environment('BAM'))
+
+        
 def suite():
     suite = unittest.TestSuite()
     suite.addTest(unittest.makeSuite(Options))
+    suite.addTest(unittest.makeSuite(Environment))
     suite.addTest(unittest.makeSuite(SessionNewTest))
     suite.addTest(unittest.makeSuite(SessionTest))
     return suite