Google Colab local runtime Google Drive upload and download using pydrive2

colab_gdrive.py

class GDrive:
  """
  Class that allows file upload and download from a 'mounted' GDrive folder
  using pydrive2 in a local Google colab runtime

  Usage:
  First create and copy your settings.yaml according to this:
  https://docs.iterative.ai/PyDrive2/oauth/
  Then copy this code as a file and import or copy/paste

  gdrive = GDrive("Machine Learning")
  gdrive.get("helper_functions.py")
  gdrive.put("helper_functions.py")
  """
  def __init__(self, mount_path):
    "Takes an authenticated pydrive2 GoogleDrive object & folder path"
    from pydrive2.auth import GoogleAuth
    from pydrive2.drive import GoogleDrive

    # Create a socat bridge between external ip address and localhost on port 8090
    !apt update && apt install socat
    !socat tcp-listen:8090,bind=172.17.0.2,reuseaddr,fork tcp:localhost:8090 &

    gauth = GoogleAuth(settings_file='settings.yaml')
    gauth.LocalWebserverAuth('localhost', [8090])
    self.drive = GoogleDrive(gauth)

    self.build_path_map()
    # get the id for the specified mount path
    self.mount_path = mount_path
    self.mount_id = self.get_path_id(mount_path)

  def build_path_map(self):
      "Creates a path map for the whole drive"
      self.path_map = {}
      #dir_list = self.drive.ListFile({'q': "mimeType = 'application/vnd.google-apps.folder' and trashed=false"}).GetList()
      f_list = self.drive.ListFile({'q': "trashed=false", "includeItemsFromAllDrives": False, "supportsAllDrives": False}).GetList()
      def get_path_from_id(id):
        for path, f in self.path_map.items():
          if f["id"] == id:
            return path
        for f in f_list:
          if f["id"] == id:
            if f["parents"][0]["isRoot"]:
              return f["title"]
            else:
              return get_path_from_id(f["parents"][0]["id"]) + "/" + f["title"]

      for f in f_list:
        # assumes only one parent
        if not f["parents"]:
          continue  # shared files
        if f["parents"][0]["isRoot"]:
          self.path_map[f["title"]] = f
        else:
          self.path_map[get_path_from_id(f["parents"][0]["id"]) + "/" + f["title"]] = f

  def get_path_id(self, path):
    "Gets a directory id matching the path or creates a new directory"
    #file_list = drive.ListFile({'q': "'root' in parents and trashed=false"}).GetList()
    try:
      return self.path_map[path]["id"]
    except KeyError:
      parent, sep, title = path.rpartition("/")
      meta = {
        "mimeType": 'application/vnd.google-apps.folder',
        "title": title,
        "parents": [{'id': self.get_path_id(parent)}],
      }
      f = self.drive.CreateFile(meta)
      f.Upload()
      self.path_map[path] = f
      return f["id"]

  def get(self, rel_filepath):
    "Gets a remote file using a path relative to the mount path"
    path = self.mount_path + "/" + rel_filepath
    f = self.path_map[path]
    # could potentially already be a GoogleDriveFile
    if isinstance(f, dict):
      file = self.drive.CreateFile(f)
    else:
      file = f
    file.GetContentFile(rel_filepath)
    self.path_map[path] = file

  def put(self, rel_filepath):
    "Uploads a local file using a path relative to the mount path"
    path = self.mount_path + "/" + rel_filepath
    parent, sep, title = path.rpartition("/")
    try:
      f = self.path_map[path]
    except KeyError:
      f = {
        "parents": [{"id": self.get_path_id(parent)}],
      }
    if isinstance(f, dict):
      file = self.drive.CreateFile(f)
    else:
      file = f
    file.SetContentFile(rel_filepath)
    file.Upload()
    self.path_map[path] = file

You will need to run the following directly in a colab/jupyter cell

!apt update && apt install socat screen
!screen -md socat tcp-listen:8090,bind=172.17.0.2,reuseaddr,fork tcp:localhost:8090