1 files changed, 66 insertions, 13 deletions

diff --git a/setup_helper.py b/setup_helper.py
index ff6b0e16..c359a16c 100644
--- a/setup_helper.py
+++ b/setup_helper.py
@@ -30,9 +30,42 @@ import distutils.archive_util
 from distutils.dir_util import mkpath
 from distutils.spawn import spawn
 
-
-def make_tarball(base_name, base_dir, compress='gzip',
-                 verbose=False, dry_run=False):
+try:
+    from pwd import getpwnam
+except ImportError:
+    getpwnam = None
+
+try:
+    from grp import getgrnam
+except ImportError:
+    getgrnam = None
+
+def _get_gid(name):
+    """Returns a gid, given a group name."""
+    if getgrnam is None or name is None:
+        return None
+    try:
+        result = getgrnam(name)
+    except KeyError:
+        result = None
+    if result is not None:
+        return result[2]
+    return None
+
+def _get_uid(name):
+    """Returns an uid, given a user name."""
+    if getpwnam is None or name is None:
+        return None
+    try:
+        result = getpwnam(name)
+    except KeyError:
+        result = None
+    if result is not None:
+        return result[2]
+    return None
+
+def make_tarball(base_name, base_dir, compress='gzip', verbose=0, dry_run=0,
+                 owner=None, group=None):
     """Create a tar file from all the files under 'base_dir'.
     This file may be compressed.
 
@@ -44,7 +77,7 @@ def make_tarball(base_name, base_dir, compress='gzip',
     For 'gzip' and 'bzip2' the internal tarfile module will be used.
     For 'compress' the .tar will be created using tarfile, and then
     we will spawn 'compress' afterwards.
-    The output tar file will be named 'base_name' + ".tar", 
+    The output tar file will be named 'base_name' + ".tar",
     possibly plus the appropriate compression extension (".gz",
     ".bz2" or ".Z").  Return the output filename.
     """
@@ -54,12 +87,14 @@ def make_tarball(base_name, base_dir, compress='gzip',
     # "create a tree of hardlinks" step!  (Would also be nice to
     # detect GNU tar to use its 'z' option and save a step.)
 
-    compress_ext = { 'gzip': ".gz",
-                     'bzip2': '.bz2',
-                     'compress': ".Z" }
+    compress_ext = {
+        'gzip': ".gz",
+        'bzip2': '.bz2',
+        'compress': ".Z",
+    }
 
     # flags for compression program, each element of list will be an argument
-    tarfile_compress_flag = {'gzip':'gz', 'bzip2':'bz2'}
+    tarfile_compress_flag = {'gzip': 'gz', 'bzip2': 'bz2'}
     compress_flags = {'compress': ["-f"]}
 
     if compress is not None and compress not in compress_ext.keys():
@@ -75,11 +110,30 @@ def make_tarball(base_name, base_dir, compress='gzip',
     mkpath(os.path.dirname(archive_name), dry_run=dry_run)
     log.info('Creating tar file %s with mode %s' % (archive_name, mode))
 
+    uid = _get_uid(owner)
+    gid = _get_gid(group)
+
+    def _set_uid_gid(tarinfo):
+        if gid is not None:
+            tarinfo.gid = gid
+            tarinfo.gname = group
+        if uid is not None:
+            tarinfo.uid = uid
+            tarinfo.uname = owner
+        return tarinfo
+
     if not dry_run:
         tar = tarfile.open(archive_name, mode=mode)
         # This recursively adds everything underneath base_dir
-        tar.add(base_dir)
-        tar.close()
+        try:
+            try:
+                # Support for the `filter' parameter was added in Python 2.7,
+                # earlier versions will raise TypeError.
+                tar.add(base_dir, filter=_set_uid_gid)
+            except TypeError:
+                tar.add(base_dir)
+        finally:
+            tar.close()
 
     if compress and compress not in tarfile_compress_flag:
         spawn([compress] + compress_flags[compress] + [archive_name],
@@ -92,11 +146,10 @@ def make_tarball(base_name, base_dir, compress='gzip',
 _custom_formats = {
     'gztar': (make_tarball, [('compress', 'gzip')], "gzip'ed tar-file"),
     'bztar': (make_tarball, [('compress', 'bzip2')], "bzip2'ed tar-file"),
-    'ztar':  (make_tarball, [('compress', 'compress')], "compressed tar file"),
-    'tar':   (make_tarball, [('compress', None)], "uncompressed tar file"),
+    'ztar': (make_tarball, [('compress', 'compress')], "compressed tar file"),
+    'tar': (make_tarball, [('compress', None)], "uncompressed tar file"),
 }
 
 # Hack in and insert ourselves into the distutils code base
 def install_custom_make_tarball():
     distutils.archive_util.ARCHIVE_FORMATS.update(_custom_formats)
-