helpers.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  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 zero_pad(data, length):
  10. """Make sure data is `length` bytes long by prepending zero bytes
  11. >>> zero_pad('foo', 5)
  12. '\\x00\\x00foo'
  13. >>> zero_pad('foo', 3)
  14. 'foo'
  15. """
  16. return '\0' * (length - len(data)) + data
  17. def exclude_path(path, patterns):
  18. """Used by create and extract sub-commands to determine
  19. if an item should be processed or not
  20. """
  21. for pattern in (patterns or []):
  22. if pattern.match(path):
  23. return isinstance(pattern, ExcludePattern)
  24. return False
  25. class IncludePattern(object):
  26. """--include PATTERN
  27. >>> py = IncludePattern('*.py')
  28. >>> foo = IncludePattern('/foo')
  29. >>> py.match('/foo/foo.py')
  30. True
  31. >>> py.match('/bar/foo.java')
  32. False
  33. >>> foo.match('/foo/foo.py')
  34. True
  35. >>> foo.match('/bar/foo.java')
  36. False
  37. >>> foo.match('/foobar/foo.py')
  38. False
  39. """
  40. def __init__(self, pattern):
  41. self.pattern = self.dirpattern = pattern
  42. if not pattern.endswith(os.path.sep):
  43. self.dirpattern += os.path.sep
  44. def match(self, path):
  45. dir, name = os.path.split(path)
  46. return (dir + os.path.sep).startswith(self.dirpattern) or fnmatchcase(name, self.pattern)
  47. def __repr__(self):
  48. return '%s(%s)' % (type(self), self.pattern)
  49. class ExcludePattern(IncludePattern):
  50. """
  51. """
  52. def walk_dir(path):
  53. st = os.lstat(path)
  54. yield path, st
  55. if stat.S_ISDIR(st.st_mode):
  56. for f in os.listdir(path):
  57. for x in walk_dir(os.path.join(path, f)):
  58. yield x
  59. def format_time(t):
  60. """Format datetime suitable for fixed length list output
  61. """
  62. if (datetime.now() - t).days < 365:
  63. return t.strftime('%b %d %H:%M')
  64. else:
  65. return t.strftime('%b %d %Y')
  66. def format_file_mode(mod):
  67. """Format file mode bits for list output
  68. """
  69. def x(v):
  70. return ''.join(v & m and s or '-'
  71. for m, s in ((4, 'r'), (2, 'w'), (1, 'x')))
  72. return '%s%s%s' % (x(mod / 64), x(mod / 8), x(mod))
  73. def format_file_size(v):
  74. """Format file size into a human friendly format
  75. """
  76. if v > 1024 * 1024 * 1024:
  77. return '%.2f GB' % (v / 1024. / 1024. / 1024.)
  78. elif v > 1024 * 1024:
  79. return '%.2f MB' % (v / 1024. / 1024.)
  80. elif v > 1024:
  81. return '%.2f kB' % (v / 1024.)
  82. else:
  83. return str(v)
  84. class IntegrityError(Exception):
  85. """
  86. """
  87. def memoize(function):
  88. cache = {}
  89. def decorated_function(*args):
  90. try:
  91. return cache[args]
  92. except KeyError:
  93. val = function(*args)
  94. cache[args] = val
  95. return val
  96. return decorated_function
  97. @memoize
  98. def uid2user(uid):
  99. try:
  100. return pwd.getpwuid(uid).pw_name
  101. except KeyError:
  102. return None
  103. @memoize
  104. def user2uid(user):
  105. try:
  106. return pwd.getpwnam(user).pw_uid
  107. except KeyError:
  108. return None
  109. @memoize
  110. def gid2group(gid):
  111. try:
  112. return grp.getgrgid(gid).gr_name
  113. except KeyError:
  114. return None
  115. @memoize
  116. def group2gid(group):
  117. try:
  118. return grp.getgrnam(group).gr_gid
  119. except KeyError:
  120. return None
  121. class Location(object):
  122. loc_re = re.compile(r'^((?:(?P<user>[^@]+)@)?(?P<host>[^:]+):)?'
  123. r'(?P<path>[^:]*)(?:::(?P<archive>[^:]+))?$')
  124. def __init__(self, text):
  125. loc = self.loc_re.match(text)
  126. loc = loc and loc.groupdict()
  127. if not loc:
  128. raise ValueError
  129. self.user = loc['user']
  130. self.host = loc['host']
  131. self.path = loc['path']
  132. if not self.host and not self.path:
  133. raise ValueError
  134. self.archive = loc['archive']
  135. def __str__(self):
  136. text = ''
  137. if self.user:
  138. text += '%s@' % self.user
  139. if self.host:
  140. text += '%s::' % self.host
  141. if self.path:
  142. text += self.path
  143. if self.archive:
  144. text += ':%s' % self.archive
  145. return text
  146. def __repr__(self):
  147. return "Location('%s')" % self
  148. def location_validator(archive=None):
  149. def validator(text):
  150. try:
  151. loc = Location(text)
  152. except ValueError:
  153. raise argparse.ArgumentTypeError('Invalid location format: "%s"' % text)
  154. if archive is True and not loc.archive:
  155. raise argparse.ArgumentTypeError('"%s": No archive specified' % text)
  156. elif archive is False and loc.archive:
  157. raise argparse.ArgumentTypeError('"%s" No archive can be specified' % text)
  158. return loc
  159. return validator