55# the BSD License: http://www.opensource.org/licenses/bsd-license.php
66
77import os
8+ import os .path
89import sys
910import select
1011import logging
1112import threading
1213import errno
1314import mmap
1415
16+ from contextlib import contextmanager
1517from subprocess import (
1618 call ,
1719 Popen ,
1820 PIPE
1921)
2022
21-
2223from .util import (
2324 LazyMixin ,
2425 stream_copy ,
@@ -223,7 +224,7 @@ class Git(LazyMixin):
223224 Set its value to 'full' to see details about the returned values.
224225 """
225226 __slots__ = ("_working_dir" , "cat_file_all" , "cat_file_header" , "_version_info" ,
226- "_git_options" )
227+ "_git_options" , "_environment" )
227228
228229 # CONFIGURATION
229230 # The size in bytes read from stdout when copying git's output to another stream
@@ -295,7 +296,8 @@ def wait(self):
295296 :raise GitCommandError: if the return status is not 0"""
296297 status = self .proc .wait ()
297298 if status != 0 :
298- raise GitCommandError (self .args , status , self .proc .stderr .read ().decode (defenc ))
299+ error_output = self .proc .stderr .read ().decode (defenc )
300+ raise GitCommandError (self .args , status , error_output )
299301 # END status handling
300302 return status
301303 # END auto interrupt
@@ -413,6 +415,9 @@ def __init__(self, working_dir=None):
413415 self ._working_dir = working_dir
414416 self ._git_options = ()
415417
418+ # Extra environment variables to pass to git commands
419+ self ._environment = {}
420+
416421 # cached command slots
417422 self .cat_file_header = None
418423 self .cat_file_all = None
@@ -536,6 +541,8 @@ def execute(self, command,
536541 # Start the process
537542 env = os .environ .copy ()
538543 env ["LC_MESSAGES" ] = "C"
544+ env .update (self ._environment )
545+
539546 proc = Popen (command ,
540547 env = env ,
541548 cwd = cwd ,
@@ -603,6 +610,73 @@ def execute(self, command,
603610 else :
604611 return stdout_value
605612
613+ def set_environment (self , ** kwargs ):
614+ """
615+ Set environment variables for future git invocations. Return all changed
616+ values in a format that can be passed back into this function to revert
617+ the changes:
618+
619+ ``Examples``::
620+
621+ old_env = self.set_environment(PWD='/tmp')
622+ self.set_environment(**old_env)
623+
624+ :param kwargs: environment variables to use for git processes
625+ :return: dict that maps environment variables to their old values
626+ """
627+ old_env = {}
628+ for key , value in kwargs .iteritems ():
629+ # set value if it is None
630+ if value is not None :
631+ if key in self ._environment :
632+ old_env [key ] = self ._environment [key ]
633+ else :
634+ old_env [key ] = None
635+ self ._environment [key ] = value
636+ # remove key from environment if its value is None
637+ elif key in self ._environment :
638+ old_env [key ] = self ._environment [key ]
639+ del self ._environment [key ]
640+ return old_env
641+
642+ @contextmanager
643+ def environment (self , ** kwargs ):
644+ """
645+ A context manager around the above set_environment to restore the
646+ environment back to its previous state after operation.
647+
648+ ``Examples``::
649+
650+ with self.environment(GIT_SSH='/bin/ssh_wrapper'):
651+ repo.remotes.origin.fetch()
652+
653+ :param kwargs: see set_environment
654+ """
655+ old_env = self .set_environment (** kwargs )
656+ try :
657+ yield
658+ finally :
659+ self .set_environment (** old_env )
660+
661+ @contextmanager
662+ def sshkey (self , sshkey_file ):
663+ """
664+ A context manager to temporarily set an SSH key for all operations that
665+ run inside it.
666+
667+ ``Examples``::
668+
669+ with self.environment(GIT_SSH=project_dir+'deployment_key'):
670+ repo.remotes.origin.fetch()
671+
672+ :param sshkey_file: Path to a private SSH key file
673+ """
674+ this_dir = os .path .dirname (__file__ )
675+ ssh_wrapper = os .path .join (this_dir , '..' , 'scripts' , 'ssh_wrapper.py' )
676+
677+ with self .environment (GIT_SSH_KEY_FILE = sshkey_file , GIT_SSH = ssh_wrapper ):
678+ yield
679+
606680 def transform_kwargs (self , split_single_char_options = False , ** kwargs ):
607681 """Transforms Python style kwargs into git command line options."""
608682 args = list ()
0 commit comments