45 automatically attach uploaded files to model output

meorg_client/cli.py

@@ -109,19 +109,31 @@ def list_endpoints():
 
 @click.command("upload")
 @click.argument("file_path", nargs=-1)
-def file_upload(file_path: tuple):
+@click.option(
+    "--attach_to",
+    default=None,
+    help="Supply a model output id to immediately attach the file to.",
+)
+def file_upload(file_path, attach_to=None):
     """
     Upload a file to the server.
 
     Prints Job ID on success, which is used by file-status to check transfer status.
+
+    If attach_to is used then no ID is returned.
     """
     client = _get_client()
 
     # Upload the file, get the job ID
-    response = _call(client.upload_files, files=list(file_path))
-    files = response.get("data").get("files")
-    for f in files:
-        click.echo(f.get("file"))
+    response = _call(client.upload_files, files=list(file_path), attach_to=attach_to)
+
+    # Different logic if we are attaching to a model output immediately
+    if not attach_to:
+        files = response.get("data").get("files")
+        for f in files:
+            click.echo(f.get("file"))
+    else:
+        click.echo("SUCCESS")
 
 
 @click.command("upload_parallel")

meorg_client/client.py

@@ -250,13 +250,16 @@ def upload_files_parallel(self, files: Union[str, Path, list], n: int = 2):
     def upload_files(
         self,
         files: Union[str, Path],
+        attach_to: str = None,
     ) -> Union[dict, requests.Response]:
         """Upload a file.
 
         Parameters
         ----------
         files : path-like, list
             Path to the file, or a list containing paths.
+        attach_to : str, optional
+            Optional model_output_id to attach the files to, by default None
 
         Returns
         -------
@@ -309,6 +312,12 @@ def upload_files(
         for fd in payload:
             fd[1][1].close()
 
+        # Automatically attach to a model output
+        if attach_to:
+            response = self.attach_files_to_model_output(
+                attach_to, files=mu.get_uploaded_file_ids(response)
+            )
+
         return response
 
     def list_files(self, id: str) -> Union[dict, requests.Response]:

meorg_client/tests/test_cli.py

@@ -89,3 +89,13 @@ def test_file_attach(runner):
     )
 
     assert result.exit_code == 0
+
+
+def test_file_upload_with_attach(runner):
+    """Test file upload with attachment via CLI."""
+    filepath = os.path.join(mu.get_installed_data_root(), "test/test.txt")
+    model_output_id = store.get("model_output_id")
+    result = runner.invoke(
+        cli.file_upload, [filepath, filepath, "--attach_to", model_output_id]
+    )
+    assert result.exit_code == 0

meorg_client/tests/test_client.py

@@ -40,6 +40,10 @@ def _get_authenticated_client() -> Client:
     return client
 
 
+def _get_test_file():
+    return os.path.join(mu.get_installed_data_root(), "test/test.txt")
+
+
 @pytest.fixture
 def client() -> Client:
     return _get_authenticated_client()
@@ -151,6 +155,13 @@ def test_upload_file_large(client: Client):
     assert client.success()
 
 
+def test_upload_files_with_attach(client: Client):
+    """Test that the upload can also attach in the same method call."""
+    filepath = _get_test_file()
+    _ = client.upload_files([filepath, filepath], attach_to=client._model_output_id)
+    assert client.success()
+
+
 def test_upload_file_parallel(client: Client, test_filepath: str):
     """Test the uploading of a file."""
     # Upload the file

meorg_client/utilities.py

@@ -84,3 +84,20 @@ def ensure_list(obj):
         The object as a list, if it is not already.
     """
     return obj if isinstance(obj, list) else [obj]
+
+
+def get_uploaded_file_ids(response):
+    """Get the file ids out of the response object.
+
+    Parameters
+    ----------
+    response : dict
+        Response dictionary from a upload call.
+
+    Returns
+    -------
+    list
+        List of file ids.
+    """
+    file_ids = [f.get("file") for f in response.get("data").get("files")]
+    return file_ids