Merge pull request #899 from sgsmob/indicator_runner

Create indicator runner main function

Author: krivard

_delphi_utils_python/delphi_utils/archive.py

@@ -26,7 +26,6 @@
 Created: 2020-08-06
 """
 
-from argparse import ArgumentParser
 from contextlib import contextmanager
 import filecmp
 from glob import glob
@@ -98,43 +97,57 @@ def diff_export_csv(
         after_df.loc[added_idx, :])
 
 
-def run_module(archive_type: str,
-               cache_dir: str,
-               export_dir: str,
-               **kwargs):
-    """Build and run an ArchiveDiffer.
+def archiver_from_params(params):
+    """Build an ArchiveDiffer from `params`.
+
+    The type of ArchiveDiffer constructed is inferred from the parameters.
 
     Parameters
     ----------
-    archive_type: str
-        Type of ArchiveDiffer to run.  Must be one of ["git", "s3"] which correspond to
-        `GitArchiveDiffer` and `S3ArchiveDiffer`, respectively.
-    cache_dir: str
-        The directory for storing most recent archived/uploaded CSVs to start diffing from.
-    export_dir: str
-        The directory with most recent exported CSVs to diff to.
-    **kwargs:
-        Keyword arguments corresponding to constructor arguments for the respective ArchiveDiffers.
+    params: Dict[str, Dict[str, Any]]
+        Dictionary of user-defined parameters with the following structure:
+        - "common":
+            - "export_dir": str, directory to which indicator output files have been exported
+        - "archive":
+            - "cache_dir": str, directory containing cached data from previous indicator runs
+            - "branch_name" (required for git archiver): str, name of git branch
+            - "override_dirty" (optional for git archiver): bool, whether to allow overwriting of
+                untracked & uncommitted changes in `cache_dir`
+            - "commit_partial_success" (optional for git archiver): bool, whether to still commit
+                even if some files were not archived and staged due to `override_dirty=False`
+            - "commit_message" (optional for git archiver): str, commit message to use
+            - "bucket_name" (required for S3 archiver): str, name of S3 bucket to which to upload
+                files
+            - "indicator_prefix" (required for S3 archiver): str, S3 prefix for files from this
+                indicator
+            - "aws_credentials" (required for S3 archiver): Dict[str, str], authentication
+                parameters for S3 to create a boto3.Session
+
+    Returns
+    -------
+    ArchiveDiffer of the inferred type.
     """
-    if archive_type == "git":
-        arch_diff = GitArchiveDiffer(cache_dir,
-                                     export_dir,
-                                     kwargs["branch_name"],
-                                     kwargs["override_dirty"],
-                                     kwargs["commit_partial_success"],
-                                     kwargs["commit_message"])
-    elif archive_type == "s3":
-        arch_diff = S3ArchiveDiffer(cache_dir,
-                                    export_dir,
-                                    kwargs["bucket_name"],
-                                    kwargs["indicator_prefix"],
-                                    kwargs["aws_credentials"])
-    elif archive_type == "filesystem":
-        arch_diff = FilesystemArchiveDiffer(cache_dir,
-                                            export_dir)
-    else:
-        raise ValueError(f"No archive type named '{archive_type}'")
-    arch_diff.run()
+    if "archive" not in params:
+        return None
+
+    # Copy to kwargs to take advantage of default arguments to archiver
+    kwargs = params["archive"]
+    kwargs["export_dir"] = params["common"]["export_dir"]
+
+    if "branch_name" in kwargs:
+        return GitArchiveDiffer(**kwargs)
+
+    if "bucket_name" in kwargs:
+        assert "indicator_prefix" in kwargs, "Missing indicator_prefix in params"
+        assert "aws_credentials" in kwargs, "Missing aws_credentials in params"
+        return S3ArchiveDiffer(**kwargs)
+
+    # Don't run the filesystem archiver if the user misspecified the archiving params
+    assert set(kwargs.keys()) == set(["cache_dir", "export_dir"]),\
+        'If you intended to run a filesystem archiver, please remove all options other than '\
+        '"cache_dir" from the "archive" params.  Otherwise, please include either "branch_name" '\
+        'or "bucket_name" to run the git or S3 archivers, respectively.'
+    return FilesystemArchiveDiffer(**kwargs)
 
 
 class ArchiveDiffer:
@@ -621,46 +634,17 @@ def update_cache(self):
         self._cache_updated = True
 
 if __name__ == "__main__":
-    parser = ArgumentParser()
-    parser.add_argument("--archive_type", required=True, type=str,
-                        choices=["git", "s3", "filesystem"],
-                        help="Type of archive differ to use.")
-    parser.add_argument("--indicator_prefix", type=str, default="",
-                        help="The prefix for S3 keys related to this indicator."
-                        " Required for `archive_type = 's3'")
-    parser.add_argument("--branch_name", type=str, default="",
-                        help=" Branch to use for `archive_type` = 'git'.")
-    parser.add_argument("--override_dirty", action="store_true",
-                        help="Whether to allow overwriting of untracked &"
-                        " uncommitted changes for `archive_type` = 'git'")
-    parser.add_argument("--commit_partial_success",  action="store_true",
-                        help="Whether to still commit for `archive_type` = "
-                        "'git' even if some files were not archived and "
-                        "staged due to `override_dirty` = False.")
-    parser.add_argument("--commit_message", type=str, default="",
-                        help="Commit message for `archive_type` = 'git'")
-    args = parser.parse_args()
-    params = read_params()
+    _params = read_params()
 
     # Autodetect whether parameters have been factored hierarchically or not
     # See https://github.com/cmu-delphi/covidcast-indicators/issues/847
     # Once all indicators have their parameters factored in to "common", "indicator", "validation",
     # and "archive", this code will be obsolete.
-    if "archive" in params:
-        archive_params = params["archive"]
-        common_params = params["common"]
-    else:
-        archive_params = params
-        common_params = params
-
-    run_module(args.archive_type,
-               archive_params["cache_dir"],
-               common_params["export_dir"],
-               aws_credentials=archive_params.get("aws_credentials", {}),
-               branch_name=args.branch_name,
-               bucket_name=archive_params.get("bucket_name", ""),
-               commit_message=args.commit_message,
-               commit_partial_success=args.commit_partial_success,
-               indicator_prefix=args.indicator_prefix,
-               override_dirty=args.override_dirty
-               )
+    #
+    # We assume that by virtue of invoking this module from the command line that the user intends
+    # to run validation.  Thus if the "archive" sub-object is not found, we interpret that to mean
+    # the parameters have not be hierarchically refactored.
+    if "archive" not in _params:
+        _params = {"archive": _params, "common": _params}
+
+    archiver_from_params(_params).run()

_delphi_utils_python/delphi_utils/runner.py

@@ -1,8 +1,11 @@
 """Indicator running utilities."""
+import argparse as ap
+import importlib
 from typing import Any, Callable, Dict, Optional
-from .archive import ArchiveDiffer
+from .archive import ArchiveDiffer, archiver_from_params
 from .utils import read_params
 from .validator.validate import Validator
+from .validator.run import validator_from_params
 
 Params = Dict[str, Any]
 
@@ -33,3 +36,14 @@ def run_indicator_pipeline(indicator_fn:  Callable[[Params], None],
         validation_report = validator.validate()
     if archiver and (not validator or validation_report.success()):
         archiver.archive()
+
+
+if __name__ == "__main__":
+    parser = ap.ArgumentParser()
+    parser.add_argument("indicator_name",
+                        type=str,
+                        help="Name of the Python package containing the indicator.  This package "
+                             "must export a `run_module(params)` function.")
+    args = parser.parse_args()
+    indicator_module = importlib.import_module(args.indicator_name)
+    run_indicator_pipeline(indicator_module.run_module, validator_from_params, archiver_from_params)

_delphi_utils_python/delphi_utils/validator/run.py

@@ -14,7 +14,7 @@ def run_module():
     validator.validate().print_and_exit()
 
 
-def from_params(params):
+def validator_from_params(params):
     """Construct a validator from `params`.
 
     Arguments

_delphi_utils_python/tests/test_archive.py

@@ -5,13 +5,15 @@
 
 from boto3 import Session
 from git import Repo, exc
+import mock
 from moto import mock_s3
 import numpy as np
 import pandas as pd
 from pandas.testing import assert_frame_equal
 import pytest
 
-from delphi_utils import ArchiveDiffer, GitArchiveDiffer, S3ArchiveDiffer
+from delphi_utils.archive import ArchiveDiffer, GitArchiveDiffer, S3ArchiveDiffer,\
+    archiver_from_params
 
 CSV_DTYPES = {"geo_id": str, "val": float, "se": float, "sample_size": float}
 
@@ -467,3 +469,143 @@ def test_run(self, tmp_path):
         assert_frame_equal(
             pd.read_csv(join(export_dir, "csv1.csv"), dtype=CSV_DTYPES),
             csv1_diff)
+
+
+class TestFromParams:
+    """Tests for creating archive differs from params."""
+
+    def test_null_creation(self):
+        """Test that a None object is created with no "archive" params."""
+        assert archiver_from_params({"common": {}}) is None
+
+    @mock.patch("delphi_utils.archive.GitArchiveDiffer")
+    def test_get_git_archiver(self, mock_archiver):
+        """Test that GitArchiveDiffer is created successfully."""
+        params = {
+            "common": {
+                "export_dir": "dir"
+            },
+            "archive": {
+                "cache_dir": "cache",
+                "branch_name": "branch",
+                "override_dirty": True,
+                "commit_partial_success": True,
+                "commit_message": "msg"
+            }
+        }
+
+        archiver_from_params(params)
+        mock_archiver.assert_called_once_with(
+            export_dir="dir",
+            cache_dir="cache",
+            branch_name="branch",
+            override_dirty=True,
+            commit_partial_success=True,
+            commit_message="msg"
+        )
+
+    @mock.patch("delphi_utils.archive.GitArchiveDiffer")
+    def test_get_git_archiver_with_defaults(self, mock_archiver):
+        """Test that GitArchiveDiffer is created successfully without optional arguments."""
+        params = {
+            "common": {
+                "export_dir": "dir"
+            },
+            "archive": {
+                "cache_dir": "cache",
+                "branch_name": "branch",
+                "commit_message": "msg"
+            }
+        }
+
+        archiver_from_params(params)
+        mock_archiver.assert_called_once_with(
+            export_dir="dir",
+            cache_dir="cache",
+            branch_name="branch",
+            commit_message="msg"
+        )
+    @mock.patch("delphi_utils.archive.S3ArchiveDiffer")
+    def test_get_s3_archiver(self, mock_archiver):
+        """Test that S3ArchiveDiffer is created successfully."""
+        params = {
+            "common": {
+                "export_dir": "dir"
+            },
+            "archive": {
+                "cache_dir": "cache",
+                "bucket_name": "bucket",
+                "indicator_prefix": "ind",
+                "aws_credentials": {"pass": "word"}
+            }
+        }
+
+        archiver_from_params(params)
+        mock_archiver.assert_called_once_with(
+            export_dir="dir",
+            cache_dir="cache",
+            bucket_name="bucket",
+            indicator_prefix="ind",
+            aws_credentials={"pass": "word"}
+        )
+
+    def test_get_s3_archiver_without_required(self):
+        """Test that S3ArchiveDiffer is not created without required arguments."""
+        params = {
+            "common": {
+                "export_dir": "dir"
+            },
+            "archive": {
+                "cache_dir": "cache",
+                "bucket_name": "bucket"
+            }
+        }
+
+        with pytest.raises(AssertionError,
+                           match="Missing indicator_prefix in params"):
+            archiver_from_params(params)
+
+        params["archive"]["indicator_prefix"] = "prefix"
+        with pytest.raises(AssertionError,
+                           match="Missing aws_credentials in params"):
+            archiver_from_params(params)
+
+    @mock.patch("delphi_utils.archive.FilesystemArchiveDiffer")
+    def test_get_filesystem_archiver(self, mock_archiver):
+        """Test that FilesystemArchiveDiffer is created successfully."""
+        params = {
+            "common": {
+                "export_dir": "dir"
+            },
+            "archive": {
+                "cache_dir": "cache"
+            }
+        }
+
+        archiver_from_params(params)
+        mock_archiver.assert_called_once_with(
+            export_dir="dir",
+            cache_dir="cache"
+        )
+
+    def test_get_filesystem_archiver_with_extra_params(self):
+        """Test that FilesystemArchiveDiffer is not created with extra parameters."""
+        params = {
+            "common": {
+                "export_dir": "dir"
+            },
+            "archive": {
+                "cache_dir": "cache",
+                "indicator_prefix": "prefix"
+            }
+        }
+
+        with pytest.raises(AssertionError,
+                           match="If you intended to run"):
+            archiver_from_params(params)
+
+        del params["archive"]["cache_dir"]
+        del params["archive"]["indicator_prefix"]
+        with pytest.raises(AssertionError,
+                           match="If you intended to run"):
+            archiver_from_params(params)