<commit_before>import base64
import json
import logging
import md5
import mimetypes
import os
import progressbar
import requests
import threading

requests_logger = logging.getLogger('requests')
requests_logger.setLevel(logging.WARNING)


class Error(Exception):
  pass


class Verb(object):
  GET = 'GET'
  PUT = 'PUT'
  DELETE = 'DELETE'


class RpcError(Error):

  def __init__(self, status, message=None, data=None):
    self.status = status
    self.message = data['error_message'] if data else message
    self.data = data

  def __str__(self):
    return self.message

  def __getitem__(self, name):
    return self.data[name]


class GoogleStorageRpcError(RpcError):
  pass


class ProgressBarThread(threading.Thread):

  def __init__(self, bar, enabled, *args, **kwargs):
    self.bar = bar
    self.enabled = enabled
    super(ProgressBarThread, self).__init__(*args, **kwargs)

  def run(self):
    super(ProgressBarThread, self).run()
    if self.enabled:
      self.bar.update(self.bar.currval + 1)


class Jetway(object):

  def __init__(self, project, name, host, secure=False, bar_enabled=True):
    if '/' not in project:
      raise ValueError('Project must be in format: <owner>/<project>')
    self.owner, self.project = project.split('/')
    self.name = name
    self.host = host
    self.scheme = 'https' if secure else 'http'
    self.gs = GoogleStorageSigner()
    self.bar_enabled = bar_enabled

  @property
  def fileset(self):
    return {
        'name': self.name,
        'project': {'owner': {'nickname': self.owner}, 'nickname': self.project},
    }

  def rpc(self, path, body=None):
    if body is None:
      body = {}
    headers = {'Content-Type': 'application/json'}
    url = '{}://{}/_api/{}'.format(self.scheme, self.host, path)
    resp = requests.post(url, data=json.dumps(body), headers=headers)
    if not (resp.status_code >= 200 and resp.status_code < 205):
      data = resp.json()
      raise RpcError(resp.status_code, data=data)
    return resp.json()

  def upload_dir(self, build_dir):
    paths_to_contents = Jetway._get_paths_to_contents_from_dir(build_dir)
    return self.write(paths_to_contents)

  def delete(self, paths):
    paths_to_contents = dict([(path, None) for path in paths])
    req = self.gs.create_sign_requests_request(Verb.DELETE, self.fileset, paths_to_contents)
    resp = self.rpc('filesets.sign_requests', req)
    return self._execute_signed_requests(resp['signed_requests'], paths_to_contents)

  def read(self, paths):
    paths_to_contents = dict([(path, None) for path in paths])
    req = self.gs.create_sign_requests_request(Verb.GET, self.fileset, paths_to_contents)
    resp = self.rpc('filesets.sign_requests', req)
    return self._execute_signed_requests(resp['signed_requests'], paths_to_contents)

  def write(self, paths_to_contents):
    req = self.gs.create_sign_requests_request(Verb.PUT, self.fileset, paths_to_contents)
    resp = self.rpc('filesets.sign_requests', req)
    return self._execute_signed_requests(resp['signed_requests'], paths_to_contents)

  def _execute_signed_requests(self, signed_requests, paths_to_contents):
    resps = {}
    errors = {}
    threads = []
    lock = threading.Lock()  # TODO(jeremydw): Thread pool.

    num_files = len(signed_requests)
    text = 'Working: %(value)d/{} (in %(elapsed)s)'
    widgets = [progressbar.FormatLabel(text.format(num_files))]
    bar = progressbar.ProgressBar(widgets=widgets, maxval=num_files)
    if self.bar_enabled:
      bar.start()

    def _execute_signed_request(req, path, content):
      with lock:
        try:
          resps[path] = self.gs.execute_signed_request(req, content)
        except GoogleStorageRpcError as e:
          errors[path] = e

    for req in signed_requests:
      path = req['path']
      thread = ProgressBarThread(bar, self.bar_enabled,
                                 target=_execute_signed_request,
                                 args=(req, path, paths_to_contents[path]))
      thread.start()
      threads.append(thread)

    for thread in threads:
      thread.join()

    if self.bar_enabled:
      bar.finish()

    return resps, errors

  @classmethod
  def _get_paths_to_contents_from_dir(cls, build_dir):
    paths_to_contents = {}
    for pre, _, files in os.walk(build_dir):
      for f in files:
        path = os.path.join(pre, f)
        fp = open(path)
        path = path.replace(build_dir, '')
        if not path.startswith('/'):
          path = '/{}'.format(path)
        content = fp.read()
        fp.close()
        if isinstance(content, unicode):
          content = content.encode('utf-8')
        paths_to_contents[path] = content
    return paths_to_contents


class GoogleStorageSigner(object):

  @staticmethod
  def create_unsigned_request(verb, path, content=None):
    req = {
      'path': path,
      'verb': verb,
    }
    if verb == Verb.PUT:
      if path.endswith('/'):
        mimetype = 'text/html'
      else:
        mimetype = mimetypes.guess_type(path)[0]
        mimetype = mimetype or 'application/octet-stream'
      md5_digest = base64.b64encode(md5.new(content).digest())
      req['headers'] = {}
      req['headers']['content_length'] = str(len(content))
      req['headers']['content_md5'] = md5_digest
      req['headers']['content_type'] = mimetype
    return req

  def create_sign_requests_request(self, verb, fileset, paths_to_contents):
    unsigned_requests = []
    for path, content in paths_to_contents.iteritems():
      req = self.create_unsigned_request(verb, path, content)
      unsigned_requests.append(req)
    return {
        'fileset': fileset,
        'unsigned_requests': unsigned_requests,
    }

  @staticmethod
  def execute_signed_request(signed_request, content=None):
    req = signed_request
    params = {
        'GoogleAccessId': req['params']['google_access_id'],
        'Signature': req['params']['signature'],
        'Expires': req['params']['expires'],
    }

    if signed_request['verb'] == Verb.PUT:
      headers = {
          'Content-Type': req['headers']['content_type'],
          'Content-MD5': req['headers']['content_md5'],
          'Content-Length': req['headers']['content_length'],
      }
      resp = requests.put(req['url'], params=params, headers=headers, data=content)

    elif signed_request['verb'] == Verb.GET:
      resp = requests.get(req['url'], params=params)

    elif signed_request['verb'] == Verb.DELETE:
      resp = requests.delete(req['url'], params=params)

    if not (resp.status_code >= 200 and resp.status_code < 205):
      raise GoogleStorageRpcError(resp.status_code, message=resp.content)

    return resp.content
<commit_msg>Move executor outside of lock.<commit_after>import base64
import json
import logging
import md5
import mimetypes
import os
import progressbar
import requests
import threading

requests_logger = logging.getLogger('requests')
requests_logger.setLevel(logging.WARNING)


class Error(Exception):
  pass


class Verb(object):
  GET = 'GET'
  PUT = 'PUT'
  DELETE = 'DELETE'


class RpcError(Error):

  def __init__(self, status, message=None, data=None):
    self.status = status
    self.message = data['error_message'] if data else message
    self.data = data

  def __str__(self):
    return self.message

  def __getitem__(self, name):
    return self.data[name]


class GoogleStorageRpcError(RpcError):
  pass


class ProgressBarThread(threading.Thread):

  def __init__(self, bar, enabled, *args, **kwargs):
    self.bar = bar
    self.enabled = enabled
    super(ProgressBarThread, self).__init__(*args, **kwargs)

  def run(self):
    super(ProgressBarThread, self).run()
    if self.enabled:
      self.bar.update(self.bar.currval + 1)


class Jetway(object):

  def __init__(self, project, name, host, secure=False, bar_enabled=True):
    if '/' not in project:
      raise ValueError('Project must be in format: <owner>/<project>')
    self.owner, self.project = project.split('/')
    self.name = name
    self.host = host
    self.scheme = 'https' if secure else 'http'
    self.gs = GoogleStorageSigner()
    self.bar_enabled = bar_enabled

  @property
  def fileset(self):
    return {
        'name': self.name,
        'project': {'owner': {'nickname': self.owner}, 'nickname': self.project},
    }

  def rpc(self, path, body=None):
    if body is None:
      body = {}
    headers = {'Content-Type': 'application/json'}
    url = '{}://{}/_api/{}'.format(self.scheme, self.host, path)
    resp = requests.post(url, data=json.dumps(body), headers=headers)
    if not (resp.status_code >= 200 and resp.status_code < 205):
      data = resp.json()
      raise RpcError(resp.status_code, data=data)
    return resp.json()

  def upload_dir(self, build_dir):
    paths_to_contents = Jetway._get_paths_to_contents_from_dir(build_dir)
    return self.write(paths_to_contents)

  def delete(self, paths):
    paths_to_contents = dict([(path, None) for path in paths])
    req = self.gs.create_sign_requests_request(Verb.DELETE, self.fileset, paths_to_contents)
    resp = self.rpc('filesets.sign_requests', req)
    return self._execute_signed_requests(resp['signed_requests'], paths_to_contents)

  def read(self, paths):
    paths_to_contents = dict([(path, None) for path in paths])
    req = self.gs.create_sign_requests_request(Verb.GET, self.fileset, paths_to_contents)
    resp = self.rpc('filesets.sign_requests', req)
    return self._execute_signed_requests(resp['signed_requests'], paths_to_contents)

  def write(self, paths_to_contents):
    req = self.gs.create_sign_requests_request(Verb.PUT, self.fileset, paths_to_contents)
    resp = self.rpc('filesets.sign_requests', req)
    return self._execute_signed_requests(resp['signed_requests'], paths_to_contents)

  def _execute_signed_requests(self, signed_requests, paths_to_contents):
    resps = {}
    errors = {}
    threads = []
    lock = threading.Lock()  # TODO(jeremydw): Thread pool.

    num_files = len(signed_requests)
    text = 'Working: %(value)d/{} (in %(elapsed)s)'
    widgets = [progressbar.FormatLabel(text.format(num_files))]
    bar = progressbar.ProgressBar(widgets=widgets, maxval=num_files)
    if self.bar_enabled:
      bar.start()

    def _execute_signed_request(req, path, content):
      error = None
      resp = None
      try:
        resp = self.gs.execute_signed_request(req, content)
      except GoogleStorageRpcError as e:
        error = e

      with lock:
        if resp is not None:
          resps[path] = resp
        if error is not None:
          errors[path] = e

    for req in signed_requests:
      path = req['path']
      thread = ProgressBarThread(bar, self.bar_enabled,
                                 target=_execute_signed_request,
                                 args=(req, path, paths_to_contents[path]))
      thread.start()
      threads.append(thread)

    for thread in threads:
      thread.join()

    if self.bar_enabled:
      bar.finish()

    return resps, errors

  @classmethod
  def _get_paths_to_contents_from_dir(cls, build_dir):
    paths_to_contents = {}
    for pre, _, files in os.walk(build_dir):
      for f in files:
        path = os.path.join(pre, f)
        fp = open(path)
        path = path.replace(build_dir, '')
        if not path.startswith('/'):
          path = '/{}'.format(path)
        content = fp.read()
        fp.close()
        if isinstance(content, unicode):
          content = content.encode('utf-8')
        paths_to_contents[path] = content
    return paths_to_contents


class GoogleStorageSigner(object):

  @staticmethod
  def create_unsigned_request(verb, path, content=None):
    req = {
      'path': path,
      'verb': verb,
    }
    if verb == Verb.PUT:
      if path.endswith('/'):
        mimetype = 'text/html'
      else:
        mimetype = mimetypes.guess_type(path)[0]
        mimetype = mimetype or 'application/octet-stream'
      md5_digest = base64.b64encode(md5.new(content).digest())
      req['headers'] = {}
      req['headers']['content_length'] = str(len(content))
      req['headers']['content_md5'] = md5_digest
      req['headers']['content_type'] = mimetype
    return req

  def create_sign_requests_request(self, verb, fileset, paths_to_contents):
    unsigned_requests = []
    for path, content in paths_to_contents.iteritems():
      req = self.create_unsigned_request(verb, path, content)
      unsigned_requests.append(req)
    return {
        'fileset': fileset,
        'unsigned_requests': unsigned_requests,
    }

  @staticmethod
  def execute_signed_request(signed_request, content=None):
    req = signed_request
    params = {
        'GoogleAccessId': req['params']['google_access_id'],
        'Signature': req['params']['signature'],
        'Expires': req['params']['expires'],
    }

    if signed_request['verb'] == Verb.PUT:
      headers = {
          'Content-Type': req['headers']['content_type'],
          'Content-MD5': req['headers']['content_md5'],
          'Content-Length': req['headers']['content_length'],
      }
      resp = requests.put(req['url'], params=params, headers=headers, data=content)

    elif signed_request['verb'] == Verb.GET:
      resp = requests.get(req['url'], params=params)

    elif signed_request['verb'] == Verb.DELETE:
      resp = requests.delete(req['url'], params=params)

    if not (resp.status_code >= 200 and resp.status_code < 205):
      raise GoogleStorageRpcError(resp.status_code, message=resp.content)

    return resp.content

{'fraction_non_alphanumeric': 0.12729833926453143, 'fraction_numerical': 0.004151838671411625, 'mean_word_length': 3.31648, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2831861', 'n_tokens_mistral': 4297, 'n_tokens_neox': 4171, 'n_words': 910}