xbmcbackup/resources/lib/pydrive/files.py

import io
import mimetypes

from googleapiclient import errors
from googleapiclient.http import MediaIoBaseUpload
from functools import wraps

from .apiattr import ApiAttribute
from .apiattr import ApiAttributeMixin
from .apiattr import ApiResource
from .apiattr import ApiResourceList
from .auth import LoadAuth


class FileNotUploadedError(RuntimeError):
  """Error trying to access metadata of file that is not uploaded."""


class ApiRequestError(IOError):
  """Error while making any API requests."""


class FileNotDownloadableError(RuntimeError):
  """Error trying to download file that is not downloadable."""


def LoadMetadata(decoratee):
  """Decorator to check if the file has metadata and fetches it if not.

  :raises: ApiRequestError, FileNotUploadedError
  """
  @wraps(decoratee)
  def _decorated(self, *args, **kwargs):
    if not self.uploaded:
      self.FetchMetadata()
    return decoratee(self, *args, **kwargs)
  return _decorated


class GoogleDriveFileList(ApiResourceList):
  """Google Drive FileList instance.

  Equivalent to Files.list() in Drive APIs.
  """

  def __init__(self, auth=None, param=None):
    """Create an instance of GoogleDriveFileList."""
    super(GoogleDriveFileList, self).__init__(auth=auth, metadata=param)

  @LoadAuth
  def _GetList(self):
    """Overwritten method which actually makes API call to list files.

    :returns: list -- list of pydrive.files.GoogleDriveFile.
    """
    self.metadata = self.auth.service.files().list(**dict(self)).execute()
    result = []
    for file_metadata in self.metadata['items']:
      tmp_file = GoogleDriveFile(
          auth=self.auth,
          metadata=file_metadata,
          uploaded=True)
      result.append(tmp_file)
    return result


class GoogleDriveFile(ApiAttributeMixin, ApiResource):
  """Google Drive File instance.

  Inherits ApiResource which inherits dict.
  Can access and modify metadata like dictionary.
  """
  content = ApiAttribute('content')
  uploaded = ApiAttribute('uploaded')
  metadata = ApiAttribute('metadata')

  def __init__(self, auth=None, metadata=None, uploaded=False):
    """Create an instance of GoogleDriveFile.

    :param auth: authorized GoogleAuth instance.
    :type auth: pydrive.auth.GoogleAuth
    :param metadata: file resource to initialize GoogleDirveFile with.
    :type metadata: dict.
    :param uploaded: True if this file is confirmed to be uploaded.
    :type uploaded: bool.
    """
    ApiAttributeMixin.__init__(self)
    ApiResource.__init__(self)
    self.metadata = {}
    self.dirty = {'content': False}
    self.auth = auth
    self.uploaded = uploaded
    if uploaded:
      self.UpdateMetadata(metadata)
    elif metadata:
      self.update(metadata)

  def __getitem__(self, key):
    """Overwrites manner of accessing Files resource.

    If this file instance is not uploaded and id is specified,
    it will try to look for metadata with Files.get().

    :param key: key of dictionary query.
    :type key: str.
    :returns: value of Files resource
    :raises: KeyError, FileNotUploadedError
    """
    try:
      return dict.__getitem__(self, key)
    except KeyError as e:
      if self.uploaded:
        raise KeyError(e)
      if self.get('id'):
        self.FetchMetadata()
        return dict.__getitem__(self, key)
      else:
        raise FileNotUploadedError()

  def SetContentString(self, content):
    """Set content of this file to be a string.

    Creates io.BytesIO instance of utf-8 encoded string.
    Sets mimeType to be 'text/plain' if not specified.

    :param content: content of the file in string.
    :type content: str.
    """
    self.content = io.BytesIO(content.encode('utf-8'))
    if self.get('mimeType') is None:
      self['mimeType'] = 'text/plain'

  def SetContentFile(self, filename):
    """Set content of this file from a file.

    Opens the file specified by this method.
    Will be read, uploaded, and closed by Upload() method.
    Sets metadata 'title' and 'mimeType' automatically if not specified.

    :param filename: name of the file to be uploaded.
    :type filename: str.
    """
    self.content = open(filename, 'rb')
   
    if self.get('title') is None:
      self['title'] = filename
    if self.get('mimeType') is None:
      self['mimeType'] = mimetypes.guess_type(filename)[0]

  def GetContentString(self):
    """Get content of this file as a string.

    :returns: str -- utf-8 decoded content of the file
    :raises: ApiRequestError, FileNotUploadedError, FileNotDownloadableError
    """
    if self.content is None or type(self.content) is not io.BytesIO:
      self.FetchContent()
    return self.content.getvalue().decode('utf-8')

  def GetContentFile(self, filename, mimetype=None):
    """Save content of this file as a local file.

    :param filename: name of the file to write to.
    :type filename: str.
    :raises: ApiRequestError, FileNotUploadedError, FileNotDownloadableError
    """
    if self.content is None or type(self.content) is not io.BytesIO:
      self.FetchContent(mimetype)
    f = open(filename, 'wb')
    f.write(self.content.getvalue())
    f.close()

  @LoadAuth
  def FetchMetadata(self):
    """Download file's metadata from id using Files.get().

    :raises: ApiRequestError, FileNotUploadedError
    """
    file_id = self.metadata.get('id') or self.get('id')
    if file_id:
      try:
        metadata = self.auth.service.files().get(fileId=file_id).execute()
      except errors.HttpError as error:
        raise ApiRequestError(error)
      else:
        self.uploaded = True
        self.UpdateMetadata(metadata)
    else:
      raise FileNotUploadedError()

  @LoadMetadata
  def FetchContent(self, mimetype=None):
    """Download file's content from download_url.

    :raises: ApiRequestError, FileNotUploadedError, FileNotDownloadableError
    """
    download_url = self.metadata.get('downloadUrl')
    if download_url:
      self.content = io.BytesIO(self._DownloadFromUrl(download_url))
      self.dirty['content'] = False
      return
    
    export_links = self.metadata.get('exportLinks')
    if export_links and export_links.get(mimetype):
      self.content = io.BytesIO(
          self._DownloadFromUrl(export_links.get(mimetype)))
      self.dirty['content'] = False
      return

    raise FileNotDownloadableError(
        'No downloadLink/exportLinks for mimetype found in metadata')

  def Upload(self, param=None):
    """Upload/update file by choosing the most efficient method.

    :param param: additional parameter to upload file.
    :type param: dict.
    :raises: ApiRequestError
    """
    if self.uploaded or self.get('id') is not None:
      if self.dirty['content']:
        self._FilesUpdate(param=param)
      else:
        self._FilesPatch(param=param)
    else:
      self._FilesInsert(param=param)

  def Delete(self):
      if self.get('id') is not None:
          self.auth.service.files().delete(fileId=self.get('id')).execute()

  @LoadAuth
  def _FilesInsert(self, param=None):
    """Upload a new file using Files.insert().

    :param param: additional parameter to upload file.
    :type param: dict.
    :raises: ApiRequestError
    """
    if param is None:
      param = {}
    param['body'] = self.GetChanges()
    try:
      if self.dirty['content']:
        param['media_body'] = self._BuildMediaBody()
      metadata = self.auth.service.files().insert(**param).execute()
    except errors.HttpError as error:
      raise ApiRequestError(error)
    else:
      self.uploaded = True
      self.dirty['content'] = False
      self.UpdateMetadata(metadata)

  @LoadAuth
  @LoadMetadata
  def _FilesUpdate(self, param=None):
    """Update metadata and/or content using Files.Update().

    :param param: additional parameter to upload file.
    :type param: dict.
    :raises: ApiRequestError, FileNotUploadedError
    """
    if param is None:
      param = {}
    param['body'] = self.GetChanges()
    param['fileId'] = self.metadata.get('id')
    try:
      if self.dirty['content']:
        param['media_body'] = self._BuildMediaBody()
      metadata = self.auth.service.files().update(**param).execute()
    except errors.HttpError as error:
      raise ApiRequestError(error)
    else:
      self.uploaded = True
      self.dirty['content'] = False
      self.UpdateMetadata(metadata)

  @LoadAuth
  @LoadMetadata
  def _FilesPatch(self, param=None):
    """Update metadata using Files.Patch().

    :param param: additional parameter to upload file.
    :type param: dict.
    :raises: ApiRequestError, FileNotUploadedError
    """
    if param is None:
      param = {}
    param['body'] = self.GetChanges()
    param['fileId'] = self.metadata.get('id')
    try:
      metadata = self.auth.service.files().patch(**param).execute()
    except errors.HttpError as error:
      raise ApiRequestError(error)
    else:
      self.UpdateMetadata(metadata)

  def _BuildMediaBody(self):
    """Build MediaIoBaseUpload to get prepared to upload content of the file.

    Sets mimeType as 'application/octet-stream' if not specified.

    :returns: MediaIoBaseUpload -- instance that will be used to upload content.
    """
    if self.get('mimeType') is None:
      self['mimeType'] = 'application/octet-stream'
   
    return MediaIoBaseUpload(self.content, self['mimeType'])

  @LoadAuth
  def _DownloadFromUrl(self, url):
    """Download file from url using provided credential.

    :param url: link of the file to download.
    :type url: str.
    :returns: str -- content of downloaded file in string.
    :raises: ApiRequestError
    """
    resp, content = self.auth.service._http.request(url)
    if resp.status != 200:
      raise ApiRequestError('Cannot download file: %s' % resp)
    return content