Repository: azavea/raster-vision

Issue:
Finish PyDocs and unit tests for utils.files
We improved error handling and added unit tests for download-related functionality in rastervision.utils.files, and now we need to do the same for upload-related functionality. Also finish adding PyDocs and unit tests for the rest of utils.files.


Discussion:

Resolving patch:
diff --git a/src/rastervision/utils/files.py b/src/rastervision/utils/files.py
--- a/src/rastervision/utils/files.py
+++ b/src/rastervision/utils/files.py
@@ -1,7 +1,6 @@
 import io
 import os
 from pathlib import Path
-import rasterio
 import shutil
 import subprocess
 import tempfile
@@ -13,7 +12,7 @@
 from google.protobuf import json_format
 
 
-class NotFoundException(Exception):
+class NotReadableError(Exception):
     pass
 
 
@@ -26,6 +25,18 @@ class ProtobufParseException(Exception):
 
 
 def make_dir(path, check_empty=False, force_empty=False, use_dirname=False):
+    """Make a directory.
+
+    Args:
+        path: path to directory
+        check_empty: if True, check that directory is empty
+        force_empty: if True, delete files if necessary to make directory
+            empty
+        use_dirname: if path is a file, use the the parent directory as path
+
+    Raises:
+        ValueError if check_empty is True and directory is not empty
+    """
     directory = path
     if use_dirname:
         directory = os.path.dirname(path)
@@ -41,8 +52,20 @@ def make_dir(path, check_empty=False, force_empty=False, use_dirname=False):
             '{} needs to be an empty directory!'.format(directory))
 
 
-def get_local_path(uri, temp_dir):
-    """Convert a URI into a corresponding local path."""
+def get_local_path(uri, download_dir):
+    """Convert a URI into a corresponding local path.
+
+    If a uri is local, return it. If it's remote, we generate a path for it
+    within download_dir. For an S3 path of form s3://<bucket>/<key>, the path
+    is <download_dir>/s3/<bucket>/<key>.
+
+    Args:
+        uri: (string) URI of file
+        download_dir: (string) path to directory
+
+    Returns:
+        (string) a local path
+    """
     if uri is None:
         return None
 
@@ -50,32 +73,68 @@ def get_local_path(uri, temp_dir):
     if parsed_uri.scheme == '':
         path = uri
     elif parsed_uri.scheme == 's3':
-        path = os.path.join(temp_dir, 's3', parsed_uri.netloc,
+        path = os.path.join(download_dir, 's3', parsed_uri.netloc,
                             parsed_uri.path[1:])
 
     return path
 
 
-def sync_dir(src_dir, dest_uri, delete=False):
-    command = ['aws', 's3', 'sync', src_dir, dest_uri]
+def sync_dir(src_dir_uri, dest_dir_uri, delete=False):
+    """Synchronize a local and remote directory.
+
+    Transfers files from source to destination directories so that the
+    destination has all the source files. If delete is True, also delete
+    files in the destination to match those in the source directory.
+
+    Args:
+        src_dir_uri: (string) URI of source directory
+        dest_dir_uri: (string) URI of destination directory
+        delete: (bool)
+    """
+    command = ['aws', 's3', 'sync', src_dir_uri, dest_dir_uri]
     if delete:
         command.append('--delete')
     subprocess.run(command)
 
 
-def _is_raster(uri, s3_test=False):
-    if s3_test:
-        uri = uri.replace('s3://', '/vsis3/')
+def start_sync(src_dir_uri, dest_dir_uri, sync_interval=600):
+    """Start syncing a directory on a schedule.
 
-    try:
-        rasterio.open(uri)
-    except Exception:
-        return False
-    return uri
+    Calls sync_dir on a schedule.
+
+    Args:
+        src_dir_uri: (string) URI of source directory
+        dest_dir_uri: (string) URI of destination directory
+        sync_interval: (int) period in seconds for syncing
+    """
+
+    def _sync_dir(delete=True):
+        sync_dir(src_dir_uri, dest_dir_uri, delete=delete)
+        thread = Timer(sync_interval, _sync_dir)
+        thread.daemon = True
+        thread.start()
+
+    if urlparse(dest_dir_uri).scheme == 's3':
+        # On first sync, we don't want to delete files on S3 to match
+        # the contents of output_dir since there's nothing there yet.
+        _sync_dir(delete=False)
 
 
 def download_if_needed(uri, download_dir):
-    """Download a file into a directory if it's remote."""
+    """Download a file into a directory if it's remote.
+
+    If uri is local, there is no need to download the file.
+
+    Args:
+        uri: (string) URI of file
+        download_dir: (string) local directory to download file into
+
+    Returns:
+        (string) path to local file
+
+    Raises:
+        NotReadableError if URI cannot be read from
+    """
     if uri is None:
         return None
 
@@ -84,25 +143,67 @@ def download_if_needed(uri, download_dir):
 
     parsed_uri = urlparse(uri)
     if parsed_uri.scheme == 's3':
-        vsis3_path = _is_raster(uri, s3_test=True)
-        if vsis3_path:
-            return vsis3_path
-
         try:
             print('Downloading {} to {}'.format(uri, path))
             s3 = boto3.client('s3')
             s3.download_file(parsed_uri.netloc, parsed_uri.path[1:], path)
         except botocore.exceptions.ClientError:
-            raise NotFoundException('Could not access {}'.format(uri))
+            raise NotReadableError('Could not read {}'.format(uri))
     else:
         not_found = not os.path.isfile(path)
         if not_found:
-            raise NotFoundException('Could not access {}'.format(uri))
+            raise NotReadableError('Could not read {}'.format(uri))
 
     return path
 
 
+def upload_if_needed(src_path, dst_uri):
+    """Upload a file if the destination is remote.
+
+    If dst_uri is local, there is no need to upload.
+
+    Args:
+        src_path: (string) path to source file
+        dst_uri: (string) URI of destination for file
+
+    Raises:
+        NotWritableError if URI cannot be written to
+    """
+    if dst_uri is None:
+        return
+
+    if not (os.path.isfile(src_path) or os.path.isdir(src_path)):
+        raise Exception('{} does not exist.'.format(src_path))
+
+    parsed_uri = urlparse(dst_uri)
+    if parsed_uri.scheme == 's3':
+        # Strip the leading slash off of the path since S3 does not expect it.
+        print('Uploading {} to {}'.format(src_path, dst_uri))
+        if os.path.isfile(src_path):
+            try:
+                s3 = boto3.client('s3')
+                s3.upload_file(src_path, parsed_uri.netloc,
+                               parsed_uri.path[1:])
+            except Exception:
+                raise NotWritableError(
+                    'Could not write {}'.format(dst_uri))
+        else:
+            sync_dir(src_path, dst_uri, delete=True)
+
+
 def file_to_str(file_uri):
+    """Download contents of text file into a string.
+
+    Args:
+        file_uri: (string) URI of file
+
+    Returns:
+        (string) with contents of text file
+
+    Raises:
+        NotReadableError if URI cannot be read from
+    """
+
     parsed_uri = urlparse(file_uri)
     if parsed_uri.scheme == 's3':
         with io.BytesIO() as file_buffer:
@@ -112,22 +213,36 @@ def file_to_str(file_uri):
                                     file_buffer)
                 return file_buffer.getvalue().decode('utf-8')
             except botocore.exceptions.ClientError:
-                raise NotFoundException('Could not access {}'.format(file_uri))
+                raise NotReadableError(
+                    'Could not read {}'.format(file_uri))
     else:
         if not os.path.isfile(file_uri):
-            raise NotFoundException('Could not access {}'.format(file_uri))
+            raise NotReadableError('Could not read {}'.format(file_uri))
         with open(file_uri, 'r') as file_buffer:
             return file_buffer.read()
 
 
 def str_to_file(content_str, file_uri):
+    """Writes string to text file.
+
+    Args:
+        content_str: string to write
+        file_uri: (string) URI of file to write
+
+    Raise:
+        NotWritableError if file_uri cannot be written
+    """
     parsed_uri = urlparse(file_uri)
     if parsed_uri.scheme == 's3':
         bucket = parsed_uri.netloc
         key = parsed_uri.path[1:]
         with io.BytesIO(bytes(content_str, encoding='utf-8')) as str_buffer:
-            s3 = boto3.client('s3')
-            s3.upload_fileobj(str_buffer, bucket, key)
+            try:
+                s3 = boto3.client('s3')
+                s3.upload_fileobj(str_buffer, bucket, key)
+            except Exception:
+                raise NotWritableError(
+                    'Could not write {}'.format(file_uri))
     else:
         make_dir(file_uri, use_dirname=True)
         with open(file_uri, 'w') as content_file:
@@ -135,6 +250,20 @@ def str_to_file(content_str, file_uri):
 
 
 def load_json_config(uri, message):
+    """Load a JSON-formatted protobuf config file.
+
+    Args:
+        uri: (string) URI of config file
+        message: (google.protobuf.message.Message) empty protobuf message of
+            to load the config into. The type needs to match the content of
+            uri.
+
+    Returns:
+        the same message passed as input with fields filled in from uri
+
+    Raises:
+        ProtobufParseException if uri cannot be parsed
+    """
     try:
         return json_format.Parse(file_to_str(uri), message)
     except json_format.ParseError:
@@ -144,42 +273,17 @@ def load_json_config(uri, message):
 
 
 def save_json_config(message, uri):
-    json_str = json_format.MessageToJson(message)
-    str_to_file(json_str, uri)
-
-
-def upload_if_needed(src_path, dst_uri):
-    """Upload file or dir if the destination is remote."""
-    if dst_uri is None:
-        return
+    """Save a protobuf object to a JSON file.
 
-    if not (os.path.isfile(src_path) or os.path.isdir(src_path)):
-        raise Exception('{} does not exist.'.format(src_path))
+    Args:
+        message: (google.protobuf.message.Message) protobuf message
+        uri: (string) URI of JSON file to write message to
 
-    parsed_uri = urlparse(dst_uri)
-    if parsed_uri.scheme == 's3':
-        # Strip the leading slash off of the path since S3 does not expect it.
-        print('Uploading {} to {}'.format(src_path, dst_uri))
-        if os.path.isfile(src_path):
-            s3 = boto3.client('s3')
-            s3.upload_file(src_path, parsed_uri.netloc, parsed_uri.path[1:])
-        else:
-            sync_dir(src_path, dst_uri, delete=True)
-
-
-def start_sync(output_dir, output_uri, sync_interval=600):
-    """Start periodically syncing a directory."""
-
-    def _sync_dir(delete=True):
-        sync_dir(output_dir, output_uri, delete=delete)
-        thread = Timer(sync_interval, _sync_dir)
-        thread.daemon = True
-        thread.start()
-
-    if urlparse(output_uri).scheme == 's3':
-        # On first sync, we don't want to delete files on S3 to match
-        # th contents of output_dir since there's nothing there yet.
-        _sync_dir(delete=False)
+    Raises:
+        NotWritableError if uri cannot be written
+    """
+    json_str = json_format.MessageToJson(message)
+    str_to_file(json_str, uri)
 
 
 # Ensure that RV temp directory exists. We need to use a custom location for


Test patch:
diff --git a/src/rastervision/label_stores/object_detection_geojson_file_test.py b/src/rastervision/label_stores/object_detection_geojson_file_test.py
--- a/src/rastervision/label_stores/object_detection_geojson_file_test.py
+++ b/src/rastervision/label_stores/object_detection_geojson_file_test.py
@@ -13,7 +13,8 @@
 from rastervision.core.crs_transformer import CRSTransformer
 from rastervision.core.box import Box
 from rastervision.core.class_map import ClassMap, ClassItem
-from rastervision.utils.files import NotFoundException, NotWritableError
+from rastervision.utils.files import (NotReadableError,
+                                      NotWritableError)
 
 
 class DoubleCRSTransformer(CRSTransformer):
@@ -171,11 +172,11 @@ def test_read_invalid_uri_readable_false(self):
                 extent=self.extent,
                 readable=False,
                 writable=False)
-        except NotFoundException:
+        except NotReadableError:
             self.fail('Should not raise exception if readable=False')
 
     def test_read_invalid_uri_readable_true(self):
-        with self.assertRaises(NotFoundException):
+        with self.assertRaises(NotReadableError):
             invalid_uri = 's3://invalid_path/invalid.json'
             ObjectDetectionGeoJSONFile(
                 invalid_uri,
@@ -259,9 +260,7 @@ def test_write_invalid_uri(self):
             extent=None,
             readable=False,
             writable=True)
-        # TODO replace with NotWritableError once
-        # files.utils upload functions are improved/tested.
-        with self.assertRaises(Exception):
+        with self.assertRaises(NotWritableError):
             label_store.save()
 
     def test_valid_uri(self):
diff --git a/src/rastervision/utils/files_test.py b/src/rastervision/utils/files_test.py
--- a/src/rastervision/utils/files_test.py
+++ b/src/rastervision/utils/files_test.py
@@ -6,52 +6,155 @@
 import boto3
 from moto import mock_s3
 
-from rastervision.utils.files import (file_to_str, NotFoundException,
-                                      load_json_config, ProtobufParseException)
+from rastervision.utils.files import (
+    file_to_str, str_to_file, download_if_needed, upload_if_needed,
+    NotReadableError, NotWritableError, load_json_config,
+    ProtobufParseException, make_dir, get_local_path)
 from rastervision.protos.machine_learning_pb2 import MachineLearning
 
 
+class TestMakeDir(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_default_args(self):
+        dir = os.path.join(self.temp_dir.name, 'hello')
+        make_dir(dir)
+        self.assertTrue(os.path.isdir(dir))
+
+    def test_check_empty(self):
+        path = os.path.join(self.temp_dir.name, 'hello', 'hello.txt')
+        dir = os.path.dirname(path)
+        str_to_file('hello', path)
+
+        make_dir(dir, check_empty=False)
+        with self.assertRaises(Exception):
+            make_dir(dir, check_empty=True)
+
+    def test_force_empty(self):
+        path = os.path.join(self.temp_dir.name, 'hello', 'hello.txt')
+        dir = os.path.dirname(path)
+        str_to_file('hello', path)
+
+        make_dir(dir, force_empty=False)
+        self.assertTrue(os.path.isfile(path))
+        make_dir(dir, force_empty=True)
+        is_empty = len(os.listdir(dir)) == 0
+        self.assertTrue(is_empty)
+
+    def test_use_dirname(self):
+        path = os.path.join(self.temp_dir.name, 'hello', 'hello.txt')
+        dir = os.path.dirname(path)
+        make_dir(path, use_dirname=True)
+        self.assertTrue(os.path.isdir(dir))
+
+
+class TestGetLocalPath(unittest.TestCase):
+    def test_local(self):
+        download_dir = '/download_dir'
+        uri = '/my/file.txt'
+        path = get_local_path(uri, download_dir)
+        self.assertEqual(path, uri)
+
+    def test_s3(self):
+        download_dir = '/download_dir'
+        uri = 's3://bucket/my/file.txt'
+        path = get_local_path(uri, download_dir)
+        self.assertEqual(path, '/download_dir/s3/bucket/my/file.txt')
+
+
 class TestFileToStr(unittest.TestCase):
+    """Test file_to_str and str_to_file."""
+
     def setUp(self):
+        # Setup mock S3 bucket.
         self.mock_s3 = mock_s3()
         self.mock_s3.start()
+        self.s3 = boto3.client('s3')
+        self.bucket_name = 'mock_bucket'
+        self.s3.create_bucket(Bucket=self.bucket_name)
 
-        # Save temp file.
+        self.content_str = 'hello'
         self.file_name = 'hello.txt'
+        self.s3_path = 's3://{}/{}'.format(self.bucket_name, self.file_name)
+
         self.temp_dir = tempfile.TemporaryDirectory()
-        self.file_path = os.path.join(self.temp_dir.name, self.file_name)
-        self.file_contents = 'hello'
-        with open(self.file_path, 'w') as myfile:
-            myfile.write(self.file_contents)
+        self.local_path = os.path.join(self.temp_dir.name, self.file_name)
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+        self.mock_s3.stop()
+
+    def test_file_to_str_local(self):
+        str_to_file(self.content_str, self.local_path)
+        content_str = file_to_str(self.local_path)
+        self.assertEqual(self.content_str, content_str)
+
+        wrong_path = '/wrongpath/x.txt'
+        with self.assertRaises(NotReadableError):
+            file_to_str(wrong_path)
+
+    def test_file_to_str_s3(self):
+        wrong_path = 's3://wrongpath/x.txt'
+
+        with self.assertRaises(NotWritableError):
+            str_to_file(self.content_str, wrong_path)
+
+        str_to_file(self.content_str, self.s3_path)
+        content_str = file_to_str(self.s3_path)
+        self.assertEqual(self.content_str, content_str)
+
+        with self.assertRaises(NotReadableError):
+            file_to_str(wrong_path)
+
+
+class TestDownloadIfNeeded(unittest.TestCase):
+    """Test download_if_needed and upload_if_needed and str_to_file."""
 
-        # Upload file to mock S3 bucket.
+    def setUp(self):
+        # Setup mock S3 bucket.
+        self.mock_s3 = mock_s3()
+        self.mock_s3.start()
         self.s3 = boto3.client('s3')
         self.bucket_name = 'mock_bucket'
         self.s3.create_bucket(Bucket=self.bucket_name)
+
+        self.content_str = 'hello'
+        self.file_name = 'hello.txt'
         self.s3_path = 's3://{}/{}'.format(self.bucket_name, self.file_name)
-        self.s3.upload_file(self.file_path, self.bucket_name, self.file_name)
+
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.local_path = os.path.join(self.temp_dir.name, self.file_name)
 
     def tearDown(self):
         self.temp_dir.cleanup()
         self.mock_s3.stop()
 
-    def test_s3(self):
-        str = file_to_str(self.s3_path)
-        self.assertEqual(str, self.file_contents)
+    def test_download_if_needed_local(self):
+        with self.assertRaises(NotReadableError):
+            download_if_needed(self.local_path, self.temp_dir.name)
 
-    def test_local(self):
-        str = file_to_str(self.file_path)
-        self.assertEqual(str, self.file_contents)
+        str_to_file(self.content_str, self.local_path)
+        upload_if_needed(self.local_path, self.local_path)
+        local_path = download_if_needed(self.local_path, self.temp_dir.name)
+        self.assertEqual(local_path, self.local_path)
 
-    def test_wrong_s3(self):
-        wrong_path = 's3://{}/{}'.format(self.bucket_name, 'x.txt')
-        with self.assertRaises(NotFoundException):
-            file_to_str(wrong_path)
+    def test_download_if_needed_s3(self):
+        with self.assertRaises(NotReadableError):
+            download_if_needed(self.s3_path, self.temp_dir.name)
 
-    def test_wrong_local(self):
-        wrong_path = '/wrongpath/x.txt'
-        with self.assertRaises(NotFoundException):
-            file_to_str(wrong_path)
+        str_to_file(self.content_str, self.local_path)
+        upload_if_needed(self.local_path, self.s3_path)
+        local_path = download_if_needed(self.s3_path, self.temp_dir.name)
+        content_str = file_to_str(local_path)
+        self.assertEqual(self.content_str, content_str)
+
+        wrong_path = 's3://wrongpath/x.txt'
+        with self.assertRaises(NotWritableError):
+            upload_if_needed(local_path, wrong_path)
 
 
 class TestLoadJsonConfig(unittest.TestCase):