Upgrade gsutil to 3.4

third_party/gsutil/gslib/wildcard_iterator.py

+# Copyright 2010 Google Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish, dis-
+# tribute, sublicense, and/or sell copies of the Software, and to permit
+# persons to whom the Software is furnished to do so, subject to the fol-
+# lowing conditions:
+#
+# The above copyright notice and this permission notice shall be included
+# in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
+# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+
+"""Implementation of wildcarding over StorageUris.
+
+StorageUri is an abstraction that Google introduced in the boto library,
+for representing storage provider-independent bucket and object names with
+a shorthand URI-like syntax (see boto/boto/storage_uri.py) The current
+class provides wildcarding support for StorageUri objects (including both
+bucket and file system objects), allowing one to express collections of
+objects with syntax like the following:
+  gs://mybucket/images/*.png
+  file:///tmp/???abc???
+
+We provide wildcarding support as part of gsutil rather than as part
+of boto because wildcarding is really part of shell command-like
+functionality.
+
+A comment about wildcard semantics: We support both single path component
+wildcards (e.g., using '*') and recursive wildcards (using '**'), for both
+file and cloud URIs. For example,
+  gs://bucket/doc/*/*.html
+would enumerate HTML files one directory down from gs://bucket/doc, while
+  gs://bucket/**/*.html
+would enumerate HTML files in all objects contained in the bucket.
+
+Note also that if you use file system wildcards it's likely your shell
+interprets the wildcarding before passing the command to gsutil. For example:
+  % gsutil cp /opt/eclipse/*/*.html gs://bucket/eclipse
+would likely be expanded by the shell into the following before running gsutil:
+  % gsutil cp /opt/eclipse/RUNNING.html gs://bucket/eclipse
+
+Note also that most shells don't support '**' wildcarding (I think only
+zsh does). If you want to use '**' wildcarding with such a shell you can
+single quote each wildcarded string, so it gets passed uninterpreted by the
+shell to gsutil (at which point gsutil will perform the wildcarding expansion):
+  % gsutil cp '/opt/eclipse/**/*.html' gs://bucket/eclipse
+"""
+
+import boto
+import fnmatch
+import glob
+import os
+import re
+import sys
+import urllib
+
+from boto.s3.prefix import Prefix
+from boto.storage_uri import BucketStorageUri
+from bucket_listing_ref import BucketListingRef
+
+# Regex to determine if a string contains any wildcards.
+WILDCARD_REGEX = re.compile('[*?\[\]]')
+
+WILDCARD_OBJECT_ITERATOR = 'wildcard_object_iterator'
+WILDCARD_BUCKET_ITERATOR = 'wildcard_bucket_iterator'
+
+
+class WildcardIterator(object):
+  """Base class for wildcarding over StorageUris.
+
+  This class implements support for iterating over StorageUris that
+  contain wildcards.
+
+  The base class is abstract; you should instantiate using the
+  wildcard_iterator() static factory method, which chooses the right
+  implementation depending on the StorageUri.
+  """
+
+  def __repr__(self):
+    """Returns string representation of WildcardIterator."""
+    return 'WildcardIterator(%s)' % self.wildcard_uri
+
+
+class CloudWildcardIterator(WildcardIterator):
+  """WildcardIterator subclass for buckets and objects.
+
+  Iterates over BucketListingRef matching the StorageUri wildcard. It's
+  much more efficient to request the Key from the BucketListingRef (via
+  GetKey()) than to request the StorageUri and then call uri.get_key()
+  to retrieve the key, for cases where you want to get metadata that's
+  available in the Bucket (for example to get the name and size of
+  each object), because that information is available in the bucket GET
+  results. If you were to iterate over URIs for such cases and then get
+  the name and size info from each resulting StorageUri, it would cause
+  an additional object GET request for each of the result URIs.
+  """
+
+  def __init__(self, wildcard_uri, proj_id_handler,
+               bucket_storage_uri_class=BucketStorageUri,
+               headers=None, debug=0):
+    """
+    Instantiates an iterator over BucketListingRef matching given wildcard URI.
+
+    Args:
+      wildcard_uri: StorageUri that contains the wildcard to iterate.
+      proj_id_handler: ProjectIdHandler to use for current command.
+      bucket_storage_uri_class: BucketStorageUri interface.
+                                Settable for testing/mocking.
+      headers: Dictionary containing optional HTTP headers to pass to boto.
+      debug: Debug level to pass in to boto connection (range 0..3).
+    """
+    self.wildcard_uri = wildcard_uri
+    # Make a copy of the headers so any updates we make during wildcard
+    # expansion aren't left in the input params (specifically, so we don't
+    # include the x-goog-project-id header needed by a subset of cases, in
+    # the data returned to caller, which could then be used in other cases
+    # where that header must not be passed).
+    if headers is None:
+      self.headers = {}
+    else:
+      self.headers = headers.copy()
+    self.proj_id_handler = proj_id_handler
+    self.bucket_storage_uri_class = bucket_storage_uri_class
+    self.debug = debug
+
+  def __iter__(self):
+    """Python iterator that gets called when iterating over cloud wildcard.
+
+    Yields:
+      BucketListingRef, or empty iterator if no matches.
+    """
+    # First handle bucket wildcarding, if any.
+    if ContainsWildcard(self.wildcard_uri.bucket_name):
+      regex = fnmatch.translate(self.wildcard_uri.bucket_name)
+      bucket_uris = []
+      prog = re.compile(regex)
+      self.proj_id_handler.FillInProjectHeaderIfNeeded(WILDCARD_BUCKET_ITERATOR,
+                                                       self.wildcard_uri,
+                                                       self.headers)
+      for b in self.wildcard_uri.get_all_buckets(headers=self.headers):
+        if prog.match(b.name):
+          # Use str(b.name) because get_all_buckets() returns Unicode
+          # string, which when used to construct x-goog-copy-src metadata
+          # requests for object-to-object copies causes pathname '/' chars
+          # to be entity-encoded (bucket%2Fdir instead of bucket/dir),
+          # which causes the request to fail.
+          uri_str = '%s://%s' % (self.wildcard_uri.scheme,
+                                 urllib.quote_plus(str(b.name)))
+          bucket_uris.append(
+              boto.storage_uri(
+                  uri_str, debug=self.debug,
+                  bucket_storage_uri_class=self.bucket_storage_uri_class,
+                  suppress_consec_slashes=False))
+    else:
+      bucket_uris = [self.wildcard_uri.clone_replace_name('')]
+
+    # Now iterate over bucket(s), and handle object wildcarding, if any.
+    self.proj_id_handler.FillInProjectHeaderIfNeeded(WILDCARD_OBJECT_ITERATOR,
+                                                     self.wildcard_uri,
+                                                     self.headers)
+    for bucket_uri in bucket_uris:
+      if self.wildcard_uri.names_bucket():
+        # Bucket-only URI.
+        yield BucketListingRef(bucket_uri, key=None, prefix=None,
+                               headers=self.headers)
+      else:
+        # URI contains an object name. If there's no wildcard just yield
+        # the needed URI.
+        if not ContainsWildcard(self.wildcard_uri.object_name):
+          uri_to_yield = bucket_uri.clone_replace_name(
+              self.wildcard_uri.object_name)
+          yield BucketListingRef(uri_to_yield, key=None, prefix=None,
+                          headers=self.headers)
+        else:
+          # URI contains a wildcard. Expand iteratively by making a prefix
+          # query of the string preceding the first wildcard char, setting
+          # delimiter=/ (unless the wildcard is **), then filtering the results
+          # by the wildcard at that level. For example given the wildcard:
+          #   gs://bucket/abc/d*e/f*.txt
+          # we would:
+          #   - get a bucket listing with prefix=abc/d, delimiter=/
+          #   - filter each result for those that start with the result + *e
+          # Assuming gs://bucket/abc/dxyze is a result from this iteration, the
+          # next iteration would:
+          #   - get a bucket listing with prefix= abc/dxyze, delimiter=/
+          #   - filter each result for those that start with the result + f.txt
+          #
+          # Initialize the iteration with bucket name from bucket_uri but
+          # object name from self.wildcard_uri. This is needed to handle cases
+          # where both the bucket and object names contain wildcards.
+          uris_needing_expansion = [
+              bucket_uri.clone_replace_name(self.wildcard_uri.object_name)]
+          while len(uris_needing_expansion) > 0:
+            uri = uris_needing_expansion.pop(0)
+            (prefix, delimiter, prefix_wildcard, suffix) = (
+                self._BuildBucketFilterStrings(uri.object_name))
+            prog = re.compile(fnmatch.translate(prefix_wildcard))
+            # List bucket for objects matching prefix up to delimiter.
+            for key in bucket_uri.get_bucket(
+                validate=False, headers=self.headers).list(
+                    prefix=prefix, delimiter=delimiter, headers=self.headers):
+              # Check that the prefix regex matches.
+              # Match rstripped key.name, to correspond with the rstripped
+              # prefix_wildcard from _BuildBucketFilterStrings.
+              if prog.match(key.name.rstrip('/')):
+                if suffix and WILDCARD_REGEX.search(suffix):
+                  # There's more wildcard left to expand.
+                  uris_needing_expansion.append(
+                      uri.clone_replace_name(key.name + suffix))
+                else:
+                  # Done expanding.
+                  if suffix:
+                    expanded_uri = uri.clone_replace_name(key.name + suffix)
+                  else:
+                    expanded_uri = uri.clone_replace_name(key.name)
+                  if isinstance(key, Prefix):
+                    yield BucketListingRef(expanded_uri, key=None, prefix=key,
+                                           headers=self.headers)
+                  else:
+                    yield BucketListingRef(expanded_uri, key=key, prefix=None,
+                                           headers=self.headers)
+
+  def _BuildBucketFilterStrings(self, wildcard):
+    """
+    Builds strings needed for querying a bucket and filtering results to
+    implement wildcard object name matching.
+
+    Args:
+      wildcard: The wildcard string to match to objects.
+
+    Returns:
+      (prefix, delimiter, prefix_wildcard, suffix)
+      where:
+        prefix is the prefix to be sent in bucket GET request.
+        delimiter is the delimiter to be sent in bucket GET request.
+        prefix_wildcard is the wildcard to be used to filter bucket GET results.
+        suffix is string to be appended to filtered bucket GET results for next
+          wildcard expansion iteration.
+
+    Raises:
+      AssertionError if wildcard doesn't contain any wildcard chars.
+    """
+    # Generate a request prefix if the object name part of the wildcard starts
+    # with a non-regex string (e.g., that's true for 'gs://bucket/abc*xyz').
+    match = WILDCARD_REGEX.search(wildcard)
+    assert match
+    if match.start() > 0:
+      # Wildcard does not occur at beginning of object name, so construct a
+      # prefix string to send to server.
+      prefix = wildcard[:match.start()]
+    else:
+      prefix = None
+    # Construct a sub-wildcard for the current path component. For
+    # example, while iterating the first prefix match result (with
+    # prefix abc/d), prefix_wildcard will be d*e and suffix would be
+    # f*.txt.
+    wildcard_part = wildcard[match.start():]
+    end = wildcard_part.find('/')
+    if end != -1:
+      wildcard_part = wildcard_part[:end+1]
+    # Remove trailing '/' so we will match gs://bucket/abc* as well as
+    # gs://bucket/abc*/ with the same wildcard regex.
+    prefix_wildcard = ((prefix or '') + wildcard_part).rstrip('/')
+    suffix = wildcard[match.end():]
+    end = suffix.find('/')
+    if end == -1:
+      suffix = ''
+    else:
+      suffix = suffix[end+1:]
+    # To implement recursive wildcarding, if prefix_wildcard suffix starts with
+    # '**' don't send a delimiter, and combine suffix at end of prefix_wildcard.
+    if prefix_wildcard.find('**') != -1:
+      delimiter = None
+      prefix_wildcard = prefix_wildcard + suffix
+      suffix = ''
+    else:
+      delimiter = '/'
+    # The following debug output is useful for tracing how the algorithm
+    # walks through a multi-part wildcard like gs://bucket/abc/d*e/f*.txt
+    if self.debug > 1:
+      sys.stderr.write(
+          'DEBUG: wildcard=%s, prefix=%s, delimiter=%s, '
+          'prefix_wildcard=%s, suffix=%s\n' %
+          (wildcard, prefix, delimiter, prefix_wildcard, suffix))
+    return (prefix, delimiter, prefix_wildcard, suffix)
+
+  def IterKeys(self):
+    """
+    Convenience iterator that runs underlying iterator and returns Key for each
+    iteration.
+
+    Yields:
+      Subclass of boto.s3.key.Key, or empty iterator if no matches.
+
+    Raises:
+      WildcardException: for bucket-only uri.
+    """
+    for bucket_listing_ref in self. __iter__():
+      if bucket_listing_ref.HasKey():
+        yield bucket_listing_ref.GetKey()
+
+  def IterUris(self):
+    """
+    Convenience iterator that runs underlying iterator and returns StorageUri
+    for each iteration.
+
+    Yields:
+      StorageUri, or empty iterator if no matches.
+    """
+    for bucket_listing_ref in self. __iter__():
+      yield bucket_listing_ref.GetUri()
+
+  def IterUrisForKeys(self):
+    """
+    Convenience iterator that runs underlying iterator and returns the
+    StorageUri for each iterated BucketListingRef that has a Key.
+
+    Yields:
+      StorageUri, or empty iterator if no matches.
+    """
+    for bucket_listing_ref in self. __iter__():
+      if bucket_listing_ref.HasKey():
+        yield bucket_listing_ref.GetUri()
+
+
+class FileWildcardIterator(WildcardIterator):
+  """WildcardIterator subclass for files and directories.
+
+  If you use recursive wildcards ('**') only a single such wildcard is
+  supported. For example you could use the wildcard '**/*.txt' to list all .txt
+  files in any subdirectory of the current directory, but you couldn't use a
+  wildcard like '**/abc/**/*.txt' (which would, if supported, let you find .txt
+  files in any subdirectory named 'abc').
+  """
+
+  def __init__(self, wildcard_uri, headers=None, debug=0):
+    """
+    Instantiate an iterator over BucketListingRefs matching given wildcard URI.
+
+    Args:
+      wildcard_uri: StorageUri that contains the wildcard to iterate.
+      headers: Dictionary containing optional HTTP headers to pass to boto.
+      debug: Debug level to pass in to boto connection (range 0..3).
+    """
+    self.wildcard_uri = wildcard_uri
+    self.headers = headers
+    self.debug = debug
+
+  def __iter__(self):
+    wildcard = self.wildcard_uri.object_name
+    match = re.search('\*\*', wildcard)
+    if match:
+      # Recursive wildcarding request ('.../**/...').
+      # Example input: wildcard = '/tmp/tmp2pQJAX/**/*'
+      base_dir = wildcard[:match.start()-1]
+      remaining_wildcard = wildcard[match.start()+2:]
+      # At this point for the above example base_dir = '/tmp/tmp2pQJAX' and
+      # remaining_wildcard = '/*'
+      if remaining_wildcard.startswith('*'):
+        raise WildcardException('Invalid wildcard with more than 2 consecutive '
+                                '*s (%s)' % wildcard)
+      # If there was no remaining wildcard past the recursive wildcard,
+      # treat it as if it were a '*'. For example, file://tmp/** is equivalent
+      # to file://tmp/**/*
+      if not remaining_wildcard:
+        remaining_wildcard = '*'
+      # Skip slash(es).
+      remaining_wildcard = remaining_wildcard.lstrip('/')
+      filepaths = []
+      for dirpath, unused_dirnames, filenames in os.walk(base_dir):
+        filepaths.extend(
+            os.path.join(dirpath, f) for f in fnmatch.filter(filenames,
+                                                             remaining_wildcard)
+        )
+    else:
+      # Not a recursive wildcarding request.
+      filepaths = glob.glob(wildcard)
+    for filepath in filepaths:
+      expanded_uri = self.wildcard_uri.clone_replace_name(filepath)
+      yield BucketListingRef(expanded_uri)
+
+  def IterKeys(self):
+    """
+    Placeholder to allow polymorphic use of WildcardIterator.
+
+    Raises:
+      WildcardException: in all cases.
+    """
+    raise WildcardException(
+        'Iterating over Keys not possible for file wildcards')
+
+  def IterUris(self):
+    """
+    Convenience iterator that runs underlying iterator and returns StorageUri
+    for each iteration.
+
+    Yields:
+      StorageUri, or empty iterator if no matches.
+    """
+    for bucket_listing_ref in self. __iter__():
+      yield bucket_listing_ref.GetUri()
+
+
+class WildcardException(StandardError):
+  """Exception thrown for invalid wildcard URIs."""
+
+  def __init__(self, reason):
+    StandardError.__init__(self)
+    self.reason = reason
+
+  def __repr__(self):
+    return 'WildcardException: %s' % self.reason
+
+  def __str__(self):
+    return 'WildcardException: %s' % self.reason
+
+
+def wildcard_iterator(uri_or_str, proj_id_handler,
+                      bucket_storage_uri_class=BucketStorageUri,
+                      headers=None, debug=0):
+  """Instantiate a WildCardIterator for the given StorageUri.
+
+  Args:
+    uri_or_str: StorageUri or URI string naming wildcard objects to iterate.
+    proj_id_handler: ProjectIdHandler to use for current command.
+    bucket_storage_uri_class: BucketStorageUri interface.
+        Settable for testing/mocking.
+    headers: Dictionary containing optional HTTP headers to pass to boto.
+    debug: Debug level to pass in to boto connection (range 0..3).
+
+  Returns:
+    A WildcardIterator that handles the requested iteration.
+  """
+
+  if isinstance(uri_or_str, basestring):
+    # Disable enforce_bucket_naming, to allow bucket names containing
+    # wildcard chars.
+    uri = boto.storage_uri(
+        uri_or_str, debug=debug, validate=False,
+        bucket_storage_uri_class=bucket_storage_uri_class,
+        suppress_consec_slashes=False)
+  else:
+    uri = uri_or_str
+
+  if uri.is_cloud_uri():
+    return CloudWildcardIterator(
+        uri, proj_id_handler,
+        bucket_storage_uri_class=bucket_storage_uri_class, headers=headers,
+        debug=debug)
+  elif uri.is_file_uri():
+    return FileWildcardIterator(uri, headers=headers, debug=debug)
+  else:
+    raise WildcardException('Unexpected type of StorageUri (%s)' % uri)
+
+
+def ContainsWildcard(uri_or_str):
+  """Checks whether uri_or_str contains a wildcard.
+
+  Args:
+    uri_or_str: StorageUri or URI string to check.
+
+  Returns:
+    bool indicator.
+  """
+  if isinstance(uri_or_str, basestring):
+    return bool(WILDCARD_REGEX.search(uri_or_str))
+  else:
+    return bool(WILDCARD_REGEX.search(uri_or_str.uri))