def _get_response_from_google_drive()

in torchdata/datapipes/iter/load/online.py [0:0]

• 25 lines of code
• 10 McCabe index (conditional complexity)

def _get_response_from_google_drive(url: str, *, timeout: Optional[float]) -> Tuple[str, StreamWrapper]:
    confirm_token = None
    with requests.Session() as session:
        if timeout is None:
            response = session.get(url, stream=True)
        else:
            response = session.get(url, timeout=timeout, stream=True)
        for k, v in response.cookies.items():
            if k.startswith("download_warning"):
                confirm_token = v
        if confirm_token is None:
            if "Quota exceeded" in str(response.content):
                raise RuntimeError(f"Google drive link {url} is currently unavailable, because the quota was exceeded.")

        if confirm_token:
            url = url + "&confirm=" + confirm_token

        if timeout is None:
            response = session.get(url, stream=True)
        else:
            response = session.get(url, timeout=timeout, stream=True)

        if "content-disposition" not in response.headers:
            raise RuntimeError("Internal error: headers don't contain content-disposition.")

        filename = re.findall('filename="(.+)"', response.headers["content-disposition"])
        if filename is None:
            raise RuntimeError("Filename could not be autodetected")
    return filename[0], StreamWrapper(response.raw)