pw_env_setup: Add deactivate steps to environment

Add deactivate steps to the environment itself and not just to the
environment folder. This means deactivate can still be run even if the
environment folder has been deleted.

Bug: 278
Change-Id: I2169d9b8c4282b97645d3c7f3385f1d9b94f0171
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/23022
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Pigweed-Auto-Submit: Rob Mohr <mohrr@google.com>
Reviewed-by: Wyatt Hepler <hepler@google.com>
diff --git a/pw_env_setup/py/pw_env_setup/env_setup.py b/pw_env_setup/py/pw_env_setup/env_setup.py
index 57ee3b5..94f4bd9 100755
--- a/pw_env_setup/py/pw_env_setup/env_setup.py
+++ b/pw_env_setup/py/pw_env_setup/env_setup.py
@@ -308,7 +308,7 @@
         self._log('')
         self._env.echo('')
 
-        self._env.hash()
+        self._env.finalize()
 
         self._env.echo(Color.bold('Sanity checking the environment:'))
         self._env.echo()
diff --git a/pw_env_setup/py/pw_env_setup/environment.py b/pw_env_setup/py/pw_env_setup/environment.py
index be49f02..3e4c9b2 100644
--- a/pw_env_setup/py/pw_env_setup/environment.py
+++ b/pw_env_setup/py/pw_env_setup/environment.py
@@ -18,6 +18,15 @@
 import os
 import re
 
+# The order here is important. On Python 2 we want StringIO.StringIO and not
+# io.StringIO. On Python 3 there is no StringIO module so we want io.StringIO.
+# Not using six because six is not a standard package we can expect to have
+# installed in the system Python.
+try:
+    from StringIO import StringIO  # type: ignore
+except ImportError:
+    from io import StringIO
+
 # Disable super() warnings since this file must be Python 2 compatible.
 # pylint: disable=super-with-arguments
 
@@ -424,6 +433,27 @@
         pass
 
 
+class Function(_Action):
+    def __init__(self, name, body, *args, **kwargs):
+        super(Function, self).__init__(*args, **kwargs)
+        self._name = name
+        self._body = body
+
+    def write(self, outs, windows=(os.name == 'nt'), replacements=()):
+        del replacements  # Unused.
+        if windows:
+            return
+
+        outs.write("""
+{name}() {{
+{body}
+}}
+        """.strip().format(name=self._name, body=self._body))
+
+    def apply(self, env):
+        pass
+
+
 class Hash(_Action):
     def write(  # pylint: disable=no-self-use
         self,
@@ -477,6 +507,7 @@
         self._allcaps = allcaps
         self._replacements = []
         self._join = Join(pathsep)
+        self._finalized = False
 
     def add_replacement(self, variable, value=None):
         self._replacements.append((variable, value))
@@ -496,31 +527,34 @@
 
     def set(self, name, value):
         """Set a variable."""
+        assert not self._finalized
         name = self.normalize_key(name)
         self._actions.append(Set(name, value))
         self._blankline()
 
     def clear(self, name):
         """Remove a variable."""
+        assert not self._finalized
         name = self.normalize_key(name)
         self._actions.append(Clear(name))
         self._blankline()
 
     def _remove(self, name, value):
         """Remove a value from a variable."""
-
+        assert not self._finalized
         name = self.normalize_key(name)
         if self.get(name, None):
             self._actions.append(Remove(name, value, self._pathsep))
 
     def remove(self, name, value):
         """Remove a value from a PATH-like variable."""
+        assert not self._finalized
         self._remove(name, value)
         self._blankline()
 
     def append(self, name, value):
         """Add a value to a PATH-like variable. Rarely used, see prepend()."""
-
+        assert not self._finalized
         name = self.normalize_key(name)
         if self.get(name, None):
             self._remove(name, value)
@@ -531,7 +565,7 @@
 
     def prepend(self, name, value):
         """Add a value to the beginning of a PATH-like variable."""
-
+        assert not self._finalized
         name = self.normalize_key(name)
         if self.get(name, None):
             self._remove(name, value)
@@ -542,30 +576,46 @@
 
     def echo(self, value='', newline=True):
         """Echo a value to the terminal."""
-
+        # echo() deliberately ignores self._finalized.
         self._actions.append(Echo(value, newline))
         if value:
             self._blankline()
 
     def comment(self, comment):
         """Add a comment to the init script."""
+        # comment() deliberately ignores self._finalized.
         self._actions.append(Comment(comment))
         self._blankline()
 
     def command(self, command, exit_on_error=True):
         """Run a command."""
-
+        # command() deliberately ignores self._finalized.
         self._actions.append(Command(command, exit_on_error=exit_on_error))
         self._blankline()
 
+    def function(self, name, body):
+        """Define a function."""
+        assert not self._finalized
+        self._actions.append(Command(name, body))
+        self._blankline()
+
     def _blankline(self):
         self._actions.append(BlankLine())
 
-    def hash(self):
-        """If required by the shell rehash the PATH variable."""
+    def finalize(self):
+        """Run cleanup at the end of environment setup."""
+        assert not self._finalized
+        self._finalized = True
         self._actions.append(Hash())
         self._blankline()
 
+        if not self._windows:
+            buf = StringIO()
+            for action in self._actions:
+                action.write_deactivate(buf, windows=self._windows)
+            self._actions.append(Function('_pw_deactivate', buf.getvalue()))
+            self._blankline()
+
     def write(self, outs):
         """Writes a shell init script to outs."""
         if self._windows:
diff --git a/pw_env_setup/util.sh b/pw_env_setup/util.sh
index 1602cad..c624c5c 100644
--- a/pw_env_setup/util.sh
+++ b/pw_env_setup/util.sh
@@ -156,6 +156,10 @@
     . "$_PW_DEACTIVATE_SH"
   fi
 
+  # If there's a _pw_deactivate function run it. Redirect output to /dev/null
+  # in case _pw_deactivate doesn't exist.
+  _pw_deactivate &> /dev/null
+
   # Restore PW_ROOT.
   PW_ROOT="$_NEW_PW_ROOT"
   export PW_ROOT