helpers.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257
  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. >>> foo.match('/foo')
  59. True
  60. """
  61. def __init__(self, pattern):
  62. self.pattern = self.dirpattern = pattern
  63. if not pattern.endswith(os.path.sep):
  64. self.dirpattern += os.path.sep
  65. def match(self, path):
  66. dir, name = os.path.split(path)
  67. return (path == self.pattern
  68. or (dir + os.path.sep).startswith(self.dirpattern)
  69. or fnmatchcase(name, self.pattern))
  70. def __repr__(self):
  71. return '%s(%s)' % (type(self), self.pattern)
  72. class ExcludePattern(IncludePattern):
  73. """
  74. """
  75. def walk_path(path, skip_inodes=None):
  76. st = os.lstat(path)
  77. if skip_inodes and (st.st_ino, st.st_dev) in skip_inodes:
  78. return
  79. yield path, st
  80. if stat.S_ISDIR(st.st_mode):
  81. for f in os.listdir(path):
  82. for x in walk_path(os.path.join(path, f), skip_inodes):
  83. yield x
  84. def format_time(t):
  85. """Format datetime suitable for fixed length list output
  86. """
  87. if (datetime.now() - t).days < 365:
  88. return t.strftime('%b %d %H:%M')
  89. else:
  90. return t.strftime('%b %d %Y')
  91. def format_file_mode(mod):
  92. """Format file mode bits for list output
  93. """
  94. def x(v):
  95. return ''.join(v & m and s or '-'
  96. for m, s in ((4, 'r'), (2, 'w'), (1, 'x')))
  97. return '%s%s%s' % (x(mod / 64), x(mod / 8), x(mod))
  98. def format_file_size(v):
  99. """Format file size into a human friendly format
  100. """
  101. if v > 1024 * 1024 * 1024:
  102. return '%.2f GB' % (v / 1024. / 1024. / 1024.)
  103. elif v > 1024 * 1024:
  104. return '%.2f MB' % (v / 1024. / 1024.)
  105. elif v > 1024:
  106. return '%.2f kB' % (v / 1024.)
  107. else:
  108. return str(v)
  109. class IntegrityError(Exception):
  110. """
  111. """
  112. def memoize(function):
  113. cache = {}
  114. def decorated_function(*args):
  115. try:
  116. return cache[args]
  117. except KeyError:
  118. val = function(*args)
  119. cache[args] = val
  120. return val
  121. return decorated_function
  122. @memoize
  123. def uid2user(uid):
  124. try:
  125. return pwd.getpwuid(uid).pw_name
  126. except KeyError:
  127. return None
  128. @memoize
  129. def user2uid(user):
  130. try:
  131. return pwd.getpwnam(user).pw_uid
  132. except KeyError:
  133. return None
  134. @memoize
  135. def gid2group(gid):
  136. try:
  137. return grp.getgrgid(gid).gr_name
  138. except KeyError:
  139. return None
  140. @memoize
  141. def group2gid(group):
  142. try:
  143. return grp.getgrnam(group).gr_gid
  144. except KeyError:
  145. return None
  146. class Location(object):
  147. """Object representing a store / archive location
  148. >>> Location('ssh://user@host:1234/some/path::archive')
  149. Location(proto='ssh', user='user', host='host', port=1234, path='/some/path', archive='archive')
  150. >>> Location('file:///some/path::archive')
  151. Location(proto='file', user=None, host=None, port=None, path='/some/path', archive='archive')
  152. >>> Location('user@host:/some/path::archive')
  153. Location(proto='ssh', user='user', host='host', port=22, path='/some/path', archive='archive')
  154. >>> Location('/some/path::archive')
  155. Location(proto='file', user=None, host=None, port=None, path='/some/path', archive='archive')
  156. """
  157. proto = user = host = port = path = archive = None
  158. ssh_re = re.compile(r'(?P<proto>ssh)://(?:(?P<user>[^@]+)@)?'
  159. r'(?P<host>[^:/#]+)(?::(?P<port>\d+))?'
  160. r'(?P<path>[^:]*)(?:::(?P<archive>.+))?')
  161. file_re = re.compile(r'(?P<proto>file)://'
  162. r'(?P<path>[^:]*)(?:::(?P<archive>.+))?')
  163. scp_re = re.compile(r'((?:(?P<user>[^@]+)@)?(?P<host>[^:/]+):)?'
  164. r'(?P<path>[^:]*)(?:::(?P<archive>.+))?')
  165. def __init__(self, text):
  166. if not self.parse(text):
  167. raise ValueError
  168. def parse(self, text):
  169. m = self.ssh_re.match(text)
  170. if m:
  171. self.proto = m.group('proto')
  172. self.user = m.group('user')
  173. self.host = m.group('host')
  174. self.port = m.group('port') and int(m.group('port')) or 22
  175. self.path = m.group('path')
  176. self.archive = m.group('archive')
  177. return True
  178. m = self.file_re.match(text)
  179. if m:
  180. self.proto = m.group('proto')
  181. self.path = m.group('path')
  182. self.archive = m.group('archive')
  183. return True
  184. m = self.scp_re.match(text)
  185. if m:
  186. self.user = m.group('user')
  187. self.host = m.group('host')
  188. self.path = m.group('path')
  189. self.archive = m.group('archive')
  190. self.proto = self.host and 'ssh' or 'file'
  191. if self.proto == 'ssh':
  192. self.port = 22
  193. return True
  194. return False
  195. def __str__(self):
  196. items = []
  197. items.append('proto=%r' % self.proto)
  198. items.append('user=%r' % self.user)
  199. items.append('host=%r' % self.host)
  200. items.append('port=%r' % self.port)
  201. items.append('path=%r'% self.path)
  202. items.append('archive=%r' % self.archive)
  203. return ', '.join(items)
  204. def __repr__(self):
  205. return "Location(%s)" % self
  206. def location_validator(archive=None):
  207. def validator(text):
  208. try:
  209. loc = Location(text)
  210. except ValueError:
  211. raise argparse.ArgumentTypeError('Invalid location format: "%s"' % text)
  212. if archive is True and not loc.archive:
  213. raise argparse.ArgumentTypeError('"%s": No archive specified' % text)
  214. elif archive is False and loc.archive:
  215. raise argparse.ArgumentTypeError('"%s" No archive can be specified' % text)
  216. return loc
  217. return validator