Merge pull request #13 from calvingiles/peek_and_bury

Add peek methods and force bury. Resolves #9 and #10.

src/shovel/__init__.py

@@ -1,3 +1,5 @@
 # -*- coding: utf-8 -*-
+from shovel.shovel import bury, dig, peek
 
-__all__ = ['dig', 'bury']
+
+__all__ = ['bury', 'dig', 'peek']

src/shovel/cli.py

@@ -15,7 +15,7 @@ def __init__(self):
                 shovel <command> [<args>]
 
                 There are two common commands provided by Shovel:
-                    shovel bury <LOCAL_DIRECTORY> <PROJECT> <DATASET> <VERSION>
+                    shovel bury <LOCAL_DIRECTORY> <PROJECT> <DATASET> <VERSION> [--force]
                         Upload a dataset to the default pit
 
                     shovel dig <LOCAL_DIRECTORY> <PROJECT> <DATASET> <VERSION>
@@ -34,15 +34,19 @@ def __init__(self):
     def bury():
         parser = argparse.ArgumentParser(
             description='Upload a dataset to the default pit',
-            usage="shovel bury local_directory project dataset",
+            usage="shovel bury local_directory project dataset [--force]",
         )
         parser.add_argument('local_directory')
         parser.add_argument('project')
         parser.add_argument('dataset')
         parser.add_argument('version')
+        parser.add_argument('--force', dest='force', action='store_true',
+                            help='upload dataset version even if it exists already')
+
         args = parser.parse_args(sys.argv[2:])
 
-        version = shovel.bury(args.local_directory, args.project, args.dataset, args.version)
+        version = shovel.bury(
+            args.local_directory, args.project, args.dataset, args.version, args.force)
         print('Created {}/{}/{}'.format(args.project, args.dataset, version))
 
     @staticmethod

src/shovel/exceptions.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+from textwrap import dedent
 
 
 class VersionError(ValueError):
@@ -19,3 +20,17 @@ class DatasetExists(ValueError):
 
 class FileExists(ValueError):
     pass
+
+
+class OutOfSync(ValueError):
+    def __init__(self, local_files, remote_files):
+        only_local = sorted(set(local_files) - set(remote_files))
+        only_remote = sorted(set(remote_files) - set(local_files))
+        local_and_remote = sorted(set(remote_files) & set(local_files))
+        msg = dedent("""
+            Local and remote not in sync.
+            Only local: {l}
+            Only remote: {r}
+            In both: {b}
+            """).strip().format(l=only_local, r=only_remote, b=local_and_remote)
+        super().__init__(msg)

src/shovel/pit.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-from shovel import exceptions, s3
+from shovel import exceptions, local, s3
 
 
 def _version_parser(version):
@@ -18,37 +18,60 @@ def _get_latest_version(versions):
 
 class Pit(object):
     def __init__(self, bucket, root):
+        if bucket is None:
+            raise ValueError("Bucket must not be None (Did you set SHOVEL_DEFAULT_BUCKET?)")
+        if root is None:
+            raise ValueError("Root must not be None (Did you set SHOVEL_DEFAULT_ROOT?)")
         self.bucket = bucket
         self.root = root
 
-    def bury(self, project, name, version, working_path):
+    def bury(self, project, name, version, working_path, force=False):
         """Upload the contents of the target path to the pit."""
         _version_parser(version)
-        if self.list_contents(project, name, version):
+        if not force and self._list_contents(project, name, version):
             raise exceptions.VersionExists
 
         s3.put_objects(working_path, self.bucket, self.root, project, name, version)
 
     def dig(self, project, name, version, working_path):
         """Download the contents of the target dataset from the pit."""
         _version_parser(version)
-        if not self.list_contents(project, name, version):
+        if not self._list_contents(project, name, version):
             raise exceptions.VersionDoesNotExist
 
         s3.get_objects(working_path, self.bucket, self.root, project, name, version)
 
-    def list_projects(self):
+    def peek(self, project=None, name=None, version=None, *, local_path=None):
+        if project is None:
+            return self._list_projects()
+        if name is None:
+            return self._list_datasets(project)
+        if version is None:
+            return self._list_versions(project, name)
+        if local_path is None:
+            return self._list_contents(project, name, version)
+
+        # if all args are provided, check they match
+        remote_files = self._list_contents(project, name, version)
+        local_files = [fname for fname, path in local.list_local(local_path)]
+
+        if sorted(local_files) != sorted(remote_files):
+            raise exceptions.OutOfSync(local_files, remote_files)
+
+        return remote_files
+
+    def _list_projects(self):
         """Return list of projects"""
         return list(s3.list_nodes(self.bucket, self.root))
 
-    def list_datasets(self, project):
+    def _list_datasets(self, project):
         """Return list of datasets for specified project"""
         return list(s3.list_nodes(self.bucket, self.root, project))
 
-    def list_versions(self, project, dataset):
+    def _list_versions(self, project, dataset):
         """Return list of versions for specified dataset"""
         return sorted(s3.list_nodes(self.bucket, self.root, project, dataset), key=_version_parser)
 
-    def list_contents(self, project, dataset, version):
+    def _list_contents(self, project, dataset, version):
         """Return list of versions for specified dataset"""
         return list(s3.list_objects(self.bucket, self.root, project, dataset, version))

src/shovel/shovel.py

@@ -8,27 +8,25 @@
 DEFAULT_ROOT = os.environ.get('SHOVEL_DEFAULT_ROOT') or 'bottomless-pit'
 
 
-def bury(project, name, version, local_path):
+def bury(project, name, version, local_path, force=False):
+    """Upload the contents of the target path to the pit."""
     pit = get_default_pit()
-    pit.bury(project, name, version, local_path)
+    pit.bury(project, name, version, local_path, force)
 
 
 def dig(project, name, version, local_path):
+    """Download the contents of the target dataset from the pit."""
     pit = get_default_pit()
     pit.dig(project, name, version, local_path)
 
 
-def peek(project, name):
+def peek(project=None, name=None, version=None, *, local_path=None):
     pit = get_default_pit()
-    versions = pit.list_versions(project, name)
-    return {
-        version: pit.list_contents(project, name, version)
-        for version in versions
-    }
+    return pit.peek(project, name, version, local_path=local_path)
 
 
 def get_default_config():
-    return Config(bucket=DEFAULT_BUCKET, root=DEFAULT_BUCKET)
+    return Config(bucket=DEFAULT_BUCKET, root=DEFAULT_ROOT)
 
 
 def get_default_pit():