Allow stdout for docs download, DRY up download code

onecodex/cli.py

@@ -175,22 +175,25 @@ def documents_upload(ctx, max_threads, files):
 
 
 @click.command('download', help='Download a file that has been shared with you')
-@click.argument('file', nargs=1, required=True)
-@click.option('--output-path', '-o', 'path', help='Path to output directory',
-              required=False, default=None, type=click.Path(exists=True))
+@click.argument('file_id', nargs=1, required=True)
+@click.option('--output-document', '-O', 'file_path', help='Write document to PATH',
+              required=False, type=click.Path(dir_okay=False, allow_dash=True))
 @click.pass_context
 @pretty_errors
 @telemetry
 @login_required
-def documents_download(ctx, file, path):
-    doc_obj = ctx.obj['API'].Documents.get(file)
+def documents_download(ctx, file_id, file_path):
+    doc_obj = ctx.obj['API'].Documents.get(file_id)
 
     if not doc_obj:
-        log.error('Could not find document {} (404 status code)'.format(file))
+        log.error('Could not find document {} (404 status code)'.format(file_id))
         ctx.exit(1)
 
-    saved_file = doc_obj.download(path=path, progressbar=True)
-    click.echo('{} saved as {}'.format(file, saved_file))
+    if file_path == '-' or file_path == '/dev/stdout':
+        doc_obj.download(file_obj=open('/dev/stdout', 'wb'), progressbar=False)
+    else:
+        path = doc_obj.download(path=file_path, progressbar=True)
+        click.echo('{} saved to {}'.format(file_id, path), err=True)
 
 
 documents.add_command(documents_list, 'list')

onecodex/models/helpers.py

@@ -1,6 +1,9 @@
+import click
 import inspect
+import os
+import requests
 
-from onecodex.exceptions import UnboundObject
+from onecodex.exceptions import OneCodexException, UnboundObject
 
 
 def as_uri(uuid, base_class):
@@ -82,3 +85,67 @@ def truncate_string(s, length=24):
         else:
             s = s + '...'
         return s
+
+
+class ResourceDownloadMixin(object):
+    def download(self, path=None, file_obj=None, progressbar=False):
+        """Downloads files from One Codex.
+
+        Parameters
+        ----------
+        path : `string`, optional
+            Full path to save the file to. If omitted, defaults to the original filename
+            in the current working directory.
+        file_obj : file-like object, optional
+            Rather than save the file to a path, write it to this file-like object.
+
+        Returns
+        -------
+        `string`
+            The path the file was downloaded to, if applicable. Otherwise, None.
+
+        Notes
+        -----
+        If no arguments specified, defaults to download the file as the original filename
+        in the current working directory. If `file_obj` given, will write data into the
+        passed file-like object. If `path` given, will download the file to the path provided,
+        but will not overwrite any existing files.
+        """
+        if path and file_obj:
+            raise OneCodexException('Please specify only one of: path, file_obj')
+
+        if path is None and file_obj is None:
+            path = os.path.join(os.getcwd(), self.filename)
+
+        if path and os.path.exists(path):
+            raise OneCodexException('{} already exists! Will not overwrite.'.format(path))
+
+        try:
+            url_data = self._resource.download_uri()
+            resp = requests.get(url_data['download_uri'], stream=True)
+
+            with (open(path, 'wb') if path else file_obj) as f_out:
+                if progressbar:
+                    with click.progressbar(length=self.size, label=self.filename) as bar:
+                        for data in resp.iter_content(chunk_size=1024):
+                            bar.update(len(data))
+                            f_out.write(data)
+                else:
+                    for data in resp.iter_content(chunk_size=1024):
+                        f_out.write(data)
+        except KeyboardInterrupt:
+            if path:
+                os.remove(path)
+        except requests.exceptions.HTTPError as exc:
+            if exc.response.status_code == 401:
+                raise OneCodexException('You must be logged in to download files.')
+            elif exc.response.status_code == 402:
+                raise OneCodexException('You must either have a premium platform account or be in '
+                                        'a notebook environment to download files.')
+            elif exc.response.status_code == 403:
+                raise OneCodexException('You are not authorized to download this file.')
+            else:
+                raise OneCodexException('Download failed with an HTTP status code {}.'.format(
+                                        exc.response.status_code))
+
+        return path

onecodex/models/misc.py

@@ -1,13 +1,10 @@
-import click
-import os
-import requests
 import six
 import warnings
 
 from onecodex.exceptions import OneCodexException
 from onecodex.lib.upload import upload_document
 from onecodex.models import OneCodexBase
-from onecodex.models.helpers import truncate_string
+from onecodex.models.helpers import truncate_string, ResourceDownloadMixin
 
 
 class Tags(OneCodexBase):
@@ -54,53 +51,9 @@ class Jobs(OneCodexBase):
     _resource_path = '/api/v1/jobs'
 
 
-class Documents(OneCodexBase):
+class Documents(OneCodexBase, ResourceDownloadMixin):
     _resource_path = '/api/v1/documents'
 
-    def download(self, path=None, progressbar=False):
-        """Downloads a document file from One Codex.
-
-        Parameters
-        ----------
-        path : string, optional
-            Full path to save the file to. If omitted, defaults to the original filename
-            in the current working directory.
-        """
-        if path is None:
-            path = os.path.join(os.getcwd(), self.filename)
-
-        if os.path.exists(path):
-            raise OneCodexException('{} already exists! Will not overwrite.'.format(path))
-
-        try:
-            url_data = self._resource.download_uri()
-            resp = requests.get(url_data['download_uri'], stream=True)
-
-            with open(path, 'wb') as f_out:
-                if progressbar:
-                    with click.progressbar(length=self.size, label=self.filename) as bar:
-                        for data in resp.iter_content(chunk_size=1024):
-                            bar.update(len(data))
-                            f_out.write(data)
-                else:
-                    for data in resp.iter_content(chunk_size=1024):
-                        f_out.write(data)
-        except KeyboardInterrupt:
-            os.remove(path)
-        except requests.exceptions.HTTPError as exc:
-            if exc.response.status_code == 401:
-                raise OneCodexException('You must be logged in to download files.')
-            elif exc.response.status_code == 402:
-                raise OneCodexException('You must either have a premium platform account or be in '
-                                        'a notebook environment to download files.')
-            elif exc.response.status_code == 403:
-                raise OneCodexException('You are not authorized to download this file.')
-            else:
-                raise OneCodexException('Download failed with an HTTP status code {}.'.format(
-                                        exc.response.status_code))
-
-        return path
-
     @classmethod
     def upload(cls, files, threads=None, log=None, progressbar=False):
         """Uploads a series of files to the One Codex server.

onecodex/models/sample.py

@@ -1,16 +1,14 @@
-import os
-import requests
 from requests.exceptions import HTTPError
 from six import string_types
 import warnings
 
 from onecodex.exceptions import OneCodexException
 from onecodex.lib.upload import upload_sequence
 from onecodex.models import OneCodexBase, Projects, Tags
-from onecodex.models.helpers import truncate_string
+from onecodex.models.helpers import truncate_string, ResourceDownloadMixin
 
 
-class Samples(OneCodexBase):
+class Samples(OneCodexBase, ResourceDownloadMixin):
     _resource_path = '/api/v1/samples'
 
     def __repr__(self):
@@ -179,36 +177,6 @@ def upload(cls, files, threads=None, metadata=None, tags=None, project=None,
         return samples
         # FIXME: pass the auth into this so we can authenticate the callback?
 
-    def download(self, path=None):
-        """
-        Downloads the original reads file (FASTA/FASTQ) from One Codex.
-
-        Note that this may only work from within a notebook session and the file
-        is not guaranteed to exist for all One Codex plan types.
-
-        Parameters
-        ----------
-        path : string, optional
-            Full path to save the file to. If omitted, defaults to the original filename
-            in the current working directory.
-        """
-        if path is None:
-            path = os.path.join(os.getcwd(), self.filename)
-        try:
-            url_data = self._resource.download_uri()
-            resp = requests.get(url_data['download_uri'], stream=True)
-            # TODO: use tqdm or ProgressBar here to display progress?
-            with open(path, 'wb') as f_out:
-                for data in resp.iter_content(chunk_size=1024):
-                    f_out.write(data)
-        except HTTPError as exc:
-            if exc.response.status_code == 402:
-                raise OneCodexException('You must either have a premium platform account or be in '
-                                        'a notebook environment to download samples.')
-            else:
-                raise OneCodexException('Download failed with an HTTP status code {}.'.format(
-                                        exc.response.status_code))
-
     def __hash__(self):
         return hash(self.id)