helpers.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253
  1. import argparse
  2. from datetime import datetime
  3. from fnmatch import fnmatchcase
  4. import grp
  5. import os
  6. import pwd
  7. import re
  8. import stat
  9. def encode_long(v):
  10. bytes = []
  11. while True:
  12. if v > 0x7f:
  13. bytes.append(0x80 | (v % 0x80))
  14. v >>= 7
  15. else:
  16. bytes.append(v)
  17. return ''.join(chr(x) for x in bytes)
  18. def decode_long(bytes):
  19. v = 0
  20. base = 0
  21. for x in bytes:
  22. b = ord(x)
  23. if b & 0x80:
  24. v += (b & 0x7f) << base
  25. base += 7
  26. else:
  27. return v + (b << base)
  28. def zero_pad(data, length):
  29. """Make sure data is `length` bytes long by prepending zero bytes
  30. >>> zero_pad('foo', 5)
  31. '\\x00\\x00foo'
  32. >>> zero_pad('foo', 3)
  33. 'foo'
  34. """
  35. return '\0' * (length - len(data)) + data
  36. def exclude_path(path, patterns):
  37. """Used by create and extract sub-commands to determine
  38. if an item should be processed or not
  39. """
  40. for pattern in (patterns or []):
  41. if pattern.match(path):
  42. return isinstance(pattern, ExcludePattern)
  43. return False
  44. class IncludePattern(object):
  45. """--include PATTERN
  46. >>> py = IncludePattern('*.py')
  47. >>> foo = IncludePattern('/foo')
  48. >>> py.match('/foo/foo.py')
  49. True
  50. >>> py.match('/bar/foo.java')
  51. False
  52. >>> foo.match('/foo/foo.py')
  53. True
  54. >>> foo.match('/bar/foo.java')
  55. False
  56. >>> foo.match('/foobar/foo.py')
  57. False
  58. """
  59. def __init__(self, pattern):
  60. self.pattern = self.dirpattern = pattern
  61. if not pattern.endswith(os.path.sep):
  62. self.dirpattern += os.path.sep
  63. def match(self, path):
  64. dir, name = os.path.split(path)
  65. return (dir + os.path.sep).startswith(self.dirpattern) or fnmatchcase(name, self.pattern)
  66. def __repr__(self):
  67. return '%s(%s)' % (type(self), self.pattern)
  68. class ExcludePattern(IncludePattern):
  69. """
  70. """
  71. def walk_path(path, skip_inodes=None):
  72. st = os.lstat(path)
  73. if skip_inodes and (st.st_ino, st.st_dev) in skip_inodes:
  74. return
  75. yield path, st
  76. if stat.S_ISDIR(st.st_mode):
  77. for f in os.listdir(path):
  78. for x in walk_path(os.path.join(path, f), skip_inodes):
  79. yield x
  80. def format_time(t):
  81. """Format datetime suitable for fixed length list output
  82. """
  83. if (datetime.now() - t).days < 365:
  84. return t.strftime('%b %d %H:%M')
  85. else:
  86. return t.strftime('%b %d %Y')
  87. def format_file_mode(mod):
  88. """Format file mode bits for list output
  89. """
  90. def x(v):
  91. return ''.join(v & m and s or '-'
  92. for m, s in ((4, 'r'), (2, 'w'), (1, 'x')))
  93. return '%s%s%s' % (x(mod / 64), x(mod / 8), x(mod))
  94. def format_file_size(v):
  95. """Format file size into a human friendly format
  96. """
  97. if v > 1024 * 1024 * 1024:
  98. return '%.2f GB' % (v / 1024. / 1024. / 1024.)
  99. elif v > 1024 * 1024:
  100. return '%.2f MB' % (v / 1024. / 1024.)
  101. elif v > 1024:
  102. return '%.2f kB' % (v / 1024.)
  103. else:
  104. return str(v)
  105. class IntegrityError(Exception):
  106. """
  107. """
  108. def memoize(function):
  109. cache = {}
  110. def decorated_function(*args):
  111. try:
  112. return cache[args]
  113. except KeyError:
  114. val = function(*args)
  115. cache[args] = val
  116. return val
  117. return decorated_function
  118. @memoize
  119. def uid2user(uid):
  120. try:
  121. return pwd.getpwuid(uid).pw_name
  122. except KeyError:
  123. return None
  124. @memoize
  125. def user2uid(user):
  126. try:
  127. return pwd.getpwnam(user).pw_uid
  128. except KeyError:
  129. return None
  130. @memoize
  131. def gid2group(gid):
  132. try:
  133. return grp.getgrgid(gid).gr_name
  134. except KeyError:
  135. return None
  136. @memoize
  137. def group2gid(group):
  138. try:
  139. return grp.getgrnam(group).gr_gid
  140. except KeyError:
  141. return None
  142. class Location(object):
  143. """Object representing a store / archive location
  144. >>> Location('ssh://user@host:1234/some/path::archive')
  145. Location(proto='ssh', user='user', host='host', port=1234, path='/some/path', archive='archive')
  146. >>> Location('file:///some/path::archive')
  147. Location(proto='file', user=None, host=None, port=None, path='/some/path', archive='archive')
  148. >>> Location('user@host:/some/path::archive')
  149. Location(proto='ssh', user='user', host='host', port=22, path='/some/path', archive='archive')
  150. >>> Location('/some/path::archive')
  151. Location(proto='file', user=None, host=None, port=None, path='/some/path', archive='archive')
  152. """
  153. proto = user = host = port = path = archive = None
  154. ssh_re = re.compile(r'(?P<proto>ssh)://(?:(?P<user>[^@]+)@)?'
  155. r'(?P<host>[^:/#]+)(?::(?P<port>\d+))?'
  156. r'(?P<path>[^:]*)(?:::(?P<archive>.+))?')
  157. file_re = re.compile(r'(?P<proto>file)://'
  158. r'(?P<path>[^:]*)(?:::(?P<archive>.+))?')
  159. scp_re = re.compile(r'((?:(?P<user>[^@]+)@)?(?P<host>[^:/]+):)?'
  160. r'(?P<path>[^:]*)(?:::(?P<archive>.+))?')
  161. def __init__(self, text):
  162. if not self.parse(text):
  163. raise ValueError
  164. def parse(self, text):
  165. m = self.ssh_re.match(text)
  166. if m:
  167. self.proto = m.group('proto')
  168. self.user = m.group('user')
  169. self.host = m.group('host')
  170. self.port = m.group('port') and int(m.group('port')) or 22
  171. self.path = m.group('path')
  172. self.archive = m.group('archive')
  173. return True
  174. m = self.file_re.match(text)
  175. if m:
  176. self.proto = m.group('proto')
  177. self.path = m.group('path')
  178. self.archive = m.group('archive')
  179. return True
  180. m = self.scp_re.match(text)
  181. if m:
  182. self.user = m.group('user')
  183. self.host = m.group('host')
  184. self.path = m.group('path')
  185. self.archive = m.group('archive')
  186. self.proto = self.host and 'ssh' or 'file'
  187. if self.proto == 'ssh':
  188. self.port = 22
  189. return True
  190. return False
  191. def __str__(self):
  192. items = []
  193. items.append('proto=%r' % self.proto)
  194. items.append('user=%r' % self.user)
  195. items.append('host=%r' % self.host)
  196. items.append('port=%r' % self.port)
  197. items.append('path=%r'% self.path)
  198. items.append('archive=%r' % self.archive)
  199. return ', '.join(items)
  200. def __repr__(self):
  201. return "Location(%s)" % self
  202. def location_validator(archive=None):
  203. def validator(text):
  204. try:
  205. loc = Location(text)
  206. except ValueError:
  207. raise argparse.ArgumentTypeError('Invalid location format: "%s"' % text)
  208. if archive is True and not loc.archive:
  209. raise argparse.ArgumentTypeError('"%s": No archive specified' % text)
  210. elif archive is False and loc.archive:
  211. raise argparse.ArgumentTypeError('"%s" No archive can be specified' % text)
  212. return loc
  213. return validator