Add 'sshkey' context manager

teeberg · teeberg · commit 261aedd2e133 · 2015-01-21T21:35:19.000-08:00
diff --git a/doc/source/tutorial.rst b/doc/source/tutorial.rst
@@ -383,6 +383,12 @@ You can easily access configuration information for a remote by accessing option
 Change configuration for a specific remote only::
     
     o.config_writer.set("pushurl", "other_url")
+
+You can also specify an SSH key to use for any operations on the remotes:
+
+    private_key_file = project_dir+'id_rsa_deployment_key'
+    with repo.git.sshkey(private_key_file):
+        o.fetch()
     
     
 Submodule Handling
diff --git a/git/cmd.py b/git/cmd.py
@@ -5,13 +5,15 @@
 # the BSD License: http://www.opensource.org/licenses/bsd-license.php
 
 import os
+import os.path
 import sys
 import select
 import logging
 import threading
 import errno
 import mmap
 
+from contextlib import contextmanager
 from subprocess import (
     call,
     Popen,
@@ -223,7 +225,7 @@ class Git(LazyMixin):
         Set its value to 'full' to see details about the returned values.
     """
     __slots__ = ("_working_dir", "cat_file_all", "cat_file_header", "_version_info",
-                 "_git_options")
+                 "_git_options", "_environment")
 
     # CONFIGURATION
     # The size in bytes read from stdout when copying git's output to another stream
@@ -413,6 +415,9 @@ def __init__(self, working_dir=None):
         self._working_dir = working_dir
         self._git_options = ()
 
+        # Extra environment variables to pass to git commands
+        self._environment = {}
+
         # cached command slots
         self.cat_file_header = None
         self.cat_file_all = None
@@ -536,6 +541,8 @@ def execute(self, command,
         # Start the process
         env = os.environ.copy()
         env["LC_MESSAGES"] = "C"
+        env.update(self._environment)
+
         proc = Popen(command,
                      env=env,
                      cwd=cwd,
@@ -608,6 +615,73 @@ def as_text(stdout_value):
         else:
             return stdout_value
 
+    def set_environment(self, **kwargs):
+        """
+        Set environment variables for future git invocations. Return all changed
+        values in a format that can be passed back into this function to revert
+        the changes:
+
+        ``Examples``::
+
+            old_env = self.set_environment(PWD='/tmp')
+            self.set_environment(**old_env)
+
+        :param kwargs: environment variables to use for git processes
+        :return: dict that maps environment variables to their old values
+        """
+        old_env = {}
+        for key, value in kwargs.iteritems():
+            # set value if it is None
+            if value is not None:
+                if key in self._environment:
+                    old_env[key] = self._environment[key]
+                else:
+                    old_env[key] = None
+                self._environment[key] = value
+            # remove key from environment if its value is None
+            elif key in self._environment:
+                old_env[key] = self._environment[key]
+                del self._environment[key]
+        return old_env
+
+    @contextmanager
+    def environment(self, **kwargs):
+        """
+        A context manager around the above set_environment to restore the
+        environment back to its previous state after operation.
+
+        ``Examples``::
+
+            with self.environment(GIT_SSH='/bin/ssh_wrapper'):
+                repo.remotes.origin.fetch()
+
+        :param kwargs: see set_environment
+        """
+        old_env = self.set_environment(**kwargs)
+        try:
+            yield
+        finally:
+            self.set_environment(**old_env)
+
+    @contextmanager
+    def sshkey(self, sshkey_file):
+        """
+        A context manager to temporarily set an SSH key for all operations that
+        run inside it.
+
+        ``Examples``::
+
+            with self.environment(GIT_SSH=project_dir+'deployment_key'):
+                repo.remotes.origin.fetch()
+
+        :param sshkey_file: Path to a private SSH key file
+        """
+        this_dir = os.path.dirname(__file__)
+        ssh_wrapper = os.path.join(this_dir, '..', 'scripts', 'ssh_wrapper.py')
+
+        with self.environment(GIT_SSH_KEY_FILE=sshkey_file, GIT_SSH=ssh_wrapper):
+            yield
+
     def transform_kwargs(self, split_single_char_options=False, **kwargs):
         """Transforms Python style kwargs into git command line options."""
         args = list()
diff --git a/scripts/ssh_wrapper.py b/scripts/ssh_wrapper.py
@@ -0,0 +1,9 @@
+#!/usr/bin/env python
+
+import os
+import subprocess
+import sys
+
+ssh_options = ['-i', os.environ['GIT_SSH_KEY_FILE']]
+ret_code = subprocess.call(['ssh'] + ssh_options + sys.argv[1:])
+sys.exit(ret_code)
