Merge pull request #3 from ASIDataScience/factor-out-publication

Factor publishing into its own file

.pre-commit-config.yaml

@@ -0,0 +1,13 @@
+repos:
+-   repo: https://github.com/pre-commit/mirrors-isort
+    rev: HEAD  # Use the revision sha / tag you want to point at
+    hooks:
+    - id: isort
+-   repo: https://github.com/ambv/black
+    rev: stable
+    hooks:
+    - id: black
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v1.2.3
+    hooks:
+    - id: flake8

faculty_mill/faculty_reporter.py

@@ -1,15 +1,11 @@
-import os
 import shutil
 from contextlib import contextmanager
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from time import sleep
-from uuid import UUID
 
 import click
-import sherlockml
-import sml.auth
 from papermill.cli import papermill
+from publish import publish
 
 from .version import print_version
 
@@ -60,7 +56,7 @@ def main(
     notebook,
     report_name,
     description=None,
-    code=False,
+    show_code=False,
     execute=True,
 ):
     """
@@ -72,10 +68,6 @@ def main(
     This command supports all flags and options that papermill supports.
     """
 
-    PROJECT_ID = UUID(os.getenv("SHERLOCKML_PROJECT_ID"))
-    USER_ID = sml.auth.user_id()
-    report_client = sherlockml.client("report")
-
     with tmpdir() as directory:
         input_path = directory / "input.ipynb"
         output_path = directory / "output.ipynb"
@@ -91,35 +83,11 @@ def main(
                 [str(input_path), str(output_path)] + click_context.args,
                 parent=click_context,
             )
-
             papermill.invoke(papermill_click_context)
         else:
             shutil.copy(input_path, output_path)
 
-        reports = {
-            report.name: report for report in report_client.list(PROJECT_ID)
-        }
-
-        if report_name in reports:
-            report_client.create_version(
-                reports[report_name].id,
-                str(output_path.relative_to("/project/")),
-                USER_ID,
-                show_code=code,
-            )
-            click.echo("Publishing report version...")
-        else:
-            report_client.create(
-                PROJECT_ID,
-                report_name,
-                str(output_path.relative_to("/project/")),
-                USER_ID,
-                show_code=code,
-            )
-            click.echo("Publishing report...")
-
-        sleep(5)
-        click.echo("Done!")
+        publish(report_name, output_path, show_code=show_code)
 
 
 if __name__ == "__main__":

faculty_mill/publish.py

@@ -0,0 +1,94 @@
+import os
+import time
+from pathlib import Path
+from typing import Optional
+from uuid import UUID
+
+import click
+import sherlockml
+
+
+def publish(
+    report_name: str,
+    path: Path,
+    show_code: bool = False,
+    report_id: Optional[UUID] = None,
+    project_id: Optional[UUID] = None,
+    user_id: Optional[UUID] = None,
+):
+    """Publish a notebook as a report.
+
+    Parameters
+    ----------
+    report_name : str
+        The name of the report
+    path : Path
+        The path of the notebook.
+    report_id : Optional[UUID], optional
+        The report ID, if you want to publish it as a version of an existing
+        report (the default is None, in which case we search for an existing
+        report with the provided name)
+    project_id : Optional[UUID], optional
+        The project ID. Only needed if not invoking from within a project.
+    user_id : Optional[UUID], optional
+        The user ID. Only needed if not invoking from within a project.
+    show_code : bool, optional
+        Whether the code should be shown in the report or not (default False)
+    """
+
+    if project_id is None:
+        project_id = UUID(os.getenv("SHERLOCKML_PROJECT_ID"))
+
+    if user_id is None:
+        user_id = sherlockml.client("account").authenticated_user_id()
+
+    report_client = sherlockml.client("report")
+
+    if report_id is None:
+        report_id = get_report_id_by_name(report_name, project_id)
+
+    if report_id is not None:
+        report_client.create_version(
+            report_id,
+            str(path.relative_to("/project/")),
+            user_id,
+            show_code=show_code,
+        )
+        click.echo("Publishing report version...")
+    else:
+        report_client.create(
+            project_id,
+            report_name,
+            str(path.relative_to("/project/")),
+            user_id,
+            show_code=show_code,
+        )
+        click.echo("Publishing report...")
+    # this is to allow farah to process the notebook before deleting it
+    time.sleep(5)
+    click.echo("Done!")
+
+
+def get_report_id_by_name(
+    report_name: str, project_id: UUID
+) -> Optional[UUID]:
+    """Get a report id if the name exists in a project
+
+    Parameters
+    ----------
+    report_name : str
+        The name of the report.
+    project_id : UUID
+        The ID of the project in which to check for the report name.
+
+    Returns
+    -------
+    Optional[UUID]
+        The report ID if a report with that name exists, else None
+    """
+
+    report_client = sherlockml.client("report")
+    reports = {
+        report.name: report.id for report in report_client.list(project_id)
+    }
+    return reports.get(report_name)

faculty_mill/version.py

@@ -1,4 +1,4 @@
-from pkg_resources import get_distribution, DistributionNotFound
+from pkg_resources import DistributionNotFound, get_distribution
 
 try:
     __version__ = get_distribution(__package__).version

setup.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from setuptools import setup, find_packages
+from setuptools import find_packages, setup
 
 requirements = ["click", "papermill", "sherlockml", "sml"]
 test_requirements = ["pytest"]

tests/test_publish.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+from unittest.mock import Mock, patch
+from uuid import uuid4
+
+from faculty_mill import publish
+
+
+def test_that_get_report_by_name_calls_the_right_endpoint():
+
+    mock_report = Mock()
+    mock_report.id = "test id"
+    mock_report.name = "test name"
+    mock_client = Mock()
+    mock_client.list.return_value = [mock_report]
+
+    with patch(
+        "sherlockml.client", return_value=mock_client
+    ) as mock_client_creator:
+
+        result = publish.get_report_id_by_name("test name", "project id")
+        mock_client_creator.assert_called_once_with("report")
+        mock_client.list.assert_called_once_with("project id")
+        assert result == mock_report.id
+
+
+@patch("sherlockml.client")
+def test_that_publish_calls_client_method_correctly_with_all_ids_set(
+    mock_client_factory
+):
+
+    mock_client = Mock()
+    mock_client_factory.return_value = mock_client
+    test_report_id = uuid4()
+    test_project_id = uuid4()
+    test_user_id = uuid4()
+    publish.publish(
+        report_name="report name",
+        path=Path("/project/test.ipynb"),
+        show_code=False,
+        report_id=test_report_id,
+        project_id=test_project_id,
+        user_id=test_user_id,
+    )