Merge pull request #32 from griembauer/timeout

Add timeout parameter to landsatxplore download

README.md

@@ -109,10 +109,11 @@ Usage: landsatxplore download [OPTIONS] [SCENES]...
   Download one or several Landsat scenes.
 
 Options:
-  -u, --username TEXT  EarthExplorer username.
-  -p, --password TEXT  EarthExplorer password.
-  -o, --output PATH    Output directory (default to current).
-  --help               Show this message and exit.
+  -u, --username TEXT    EarthExplorer username.
+  -p, --password TEXT    EarthExplorer password.
+  -o, --output PATH      Output directory (default to current).
+  -t, --timeout INTEGER  Download timeout in seconds (default 300s).
+  --help                 Show this message and exit.
 ```
 
 ## API

README.rst

@@ -73,10 +73,12 @@ Downloading
       Download one or several Landsat scenes.
 
     Options:
-      -u, --username TEXT  EarthExplorer username.
-      -p, --password TEXT  EarthExplorer password.
-      -o, --output PATH    Output directory.
-      --help               Show this message and exit.
+      -u, --username TEXT    EarthExplorer username.
+      -p, --password TEXT    EarthExplorer password.
+      -o, --output PATH      Output directory (default to current).
+      -t, --timeout INTEGER  Download timeout in seconds (default 300s).
+      --help                 Show this message and exit.
+
 
 API
 ---

landsatxplore/cli.py

@@ -90,15 +90,17 @@ def search(username, password, dataset, location, bbox, clouds, start, end, outp
               envvar='LANDSATXPLORE_PASSWORD')
 @click.option('--output', '-o', type=click.Path(exists=True, dir_okay=True),
               default='.', help='Output directory.')
+@click.option('--timeout', '-t', type=click.INT, default=300,
+              help='Download timeout in seconds.')
 @click.argument('scenes', type=click.STRING, nargs=-1)
-def download(username, password, output, scenes):
+def download(username, password, output, timeout, scenes):
     """Download one or several Landsat scenes."""
     ee = EarthExplorer(username, password)
     output_dir = os.path.abspath(output)
     for scene in scenes:
         if not ee.logged_in():
             ee = EarthExplorer(username, password)
-        ee.download(scene, output_dir)
+        ee.download(scene, output_dir, timeout)
     ee.logout()
 
 

landsatxplore/earthexplorer.py

@@ -70,22 +70,26 @@ def logout(self):
         """Log out from Earth Explorer."""
         self.session.get(EE_LOGOUT_URL)
 
-    def _download(self, url, output_dir, chunk_size=1024):
+    def _download(self, url, output_dir, timeout, chunk_size=1024):
         """Download remote file given its URL."""
-        with self.session.get(url, stream=True, allow_redirects=True) as r:
-            file_size = int(r.headers.get("Content-Length"))
-            with tqdm(total=file_size, unit_scale=True, unit='B', unit_divisor=1024) as pbar:
-                local_filename = r.headers['Content-Disposition'].split('=')[-1]
-                local_filename = local_filename.replace("\"", "")
-                local_filename = os.path.join(output_dir, local_filename)
-                with open(local_filename, 'wb') as f:
-                    for chunk in r.iter_content(chunk_size=chunk_size):
-                        if chunk:
-                            f.write(chunk)
-                            pbar.update(chunk_size)
+        try:
+            with self.session.get(url, stream=True, allow_redirects=True, timeout=timeout) as r:
+                file_size = int(r.headers.get("Content-Length"))
+                with tqdm(total=file_size, unit_scale=True, unit='B', unit_divisor=1024) as pbar:
+                    local_filename = r.headers['Content-Disposition'].split('=')[-1]
+                    local_filename = local_filename.replace("\"", "")
+                    local_filename = os.path.join(output_dir, local_filename)
+                    with open(local_filename, 'wb') as f:
+                        for chunk in r.iter_content(chunk_size=chunk_size):
+                            if chunk:
+                                f.write(chunk)
+                                pbar.update(chunk_size)
+        except requests.exceptions.Timeout:
+            raise EarthExplorerError(
+                'Connection timeout after {} seconds.'.format(timeout))
         return local_filename
 
-    def download(self, scene_id, output_dir):
+    def download(self, scene_id, output_dir, timeout=300):
         """Download a Landsat scene given its identifier and an output
         directory.
         """
@@ -94,5 +98,5 @@ def download(self, scene_id, output_dir):
         if is_product_id(scene_id):
             scene_id = self.api.lookup(dataset, [scene_id], inverse=True)[0]
         url = EE_DOWNLOAD_URL.format(dataset_id=DATASETS[dataset], scene_id=scene_id)
-        filename = self._download(url, output_dir)
+        filename = self._download(url, output_dir, timeout=timeout)
         return filename