Merge pull request #35 from kartoza/feature/435-download-zip

feature/435: Download zip layer

django_project/cplus_api/models/layer.py

@@ -1,5 +1,6 @@
 import os
 import uuid
+from zipfile import ZipFile
 from django.db import models
 from django.utils.translation import gettext_lazy as _
 from django.conf import settings
@@ -113,16 +114,16 @@ class PrivacyTypes(models.TextChoices):
         blank=True
     )
 
-    def download_to_working_directory(self, base_dir):
+    def download_to_working_directory(self, base_dir: str):
         if not self.is_available():
             return None
-        dir_path = os.path.join(
+        dir_path: str = os.path.join(
             base_dir,
             self.component_type
         )
         if not os.path.exists(dir_path):
             os.makedirs(dir_path)
-        file_path = os.path.join(
+        file_path: str = os.path.join(
             dir_path,
             os.path.basename(self.file.name)
         )
@@ -131,6 +132,21 @@ def download_to_working_directory(self, base_dir):
                 destination.write(chunk)
         self.last_used_on = timezone.now()
         self.save(update_fields=['last_used_on'])
+        if file_path.endswith('.zip'):
+            extract_path = os.path.join(
+                dir_path,
+                os.path.basename(file_path).replace('.zip', '_zip')
+            )
+            with ZipFile(file_path, 'r') as zip_ref:
+                zip_ref.extractall(extract_path)
+            shapefile = [
+                file for file in os.listdir(extract_path)
+                if file.endswith('.shp')
+            ]
+            if shapefile:
+                return os.path.join(extract_path, shapefile[0])
+            else:
+                return None
         return file_path
 
     def is_available(self):

django_project/cplus_api/models/scenario.py

@@ -9,7 +9,9 @@
 DEFAULT_BASE_DIR = '/home/web/media'
 EXCLUDED_OUTPUT_DIR_NAMES = [
     'ncs_carbon', 'ncs_pathway', 'priority_layer',
-    'ncs_carbons', 'ncs_pathways', 'priority_layers'
+    'ncs_carbons', 'ncs_pathways', 'priority_layers',
+    'mask_layer', 'sieve_mask_layer', 'snap_layer',
+    'mask_layers', 'sieve_mask_layers', 'snap_layers'
 ]
 
 

django_project/cplus_api/tests/test_model_layer.py

@@ -1,13 +1,17 @@
-from django.test import TestCase
+import os
+import tempfile
+from django.contrib.auth.models import User
+from core.settings.utils import absolute_path
 from cplus_api.tests.factories import InputLayerF, OutputLayerF
 from cplus_api.models.layer import (
     input_layer_dir_path,
     output_layer_dir_path,
-    InputLayer
+    InputLayer,
 )
+from cplus_api.tests.common import BaseInitData
 
 
-class TestModelLayer(TestCase):
+class TestModelLayer(BaseInitData):
 
     def test_input_layer_dir_path(self):
         # private layer
@@ -55,3 +59,74 @@ def test_output_layer_dir_path(self):
             f'{str(output_layer.owner.pk)}/{str(output_layer.scenario.uuid)}/'
             'test.tif'
         )
+
+    def test_download_to_working_directory(self):
+        owner = User.objects.first()
+
+        # Test tif
+        input_layer = InputLayerF.create(owner=owner)
+        file_path = absolute_path(
+            'cplus_api', 'tests', 'data',
+            'models', 'test_model_1.tif'
+        )
+        self.store_layer_file(input_layer, file_path)
+        tmp_dir = tempfile.mkdtemp()
+        file_path = input_layer.download_to_working_directory(tmp_dir)
+        self.assertTrue(file_path.startswith(
+            os.path.join(tmp_dir, input_layer.get_component_type_display())
+        ))
+        self.assertTrue(os.path.exists(file_path))
+
+        # test zip
+        input_layer_2 = InputLayerF.create(
+            owner=owner,
+            layer_type=InputLayer.LayerTypes.VECTOR,
+            component_type=InputLayer.ComponentTypes.MASK_LAYER
+        )
+        file_path = absolute_path(
+            'cplus_api', 'tests', 'data',
+            'mask_layers', 'shapefile.zip'
+        )
+        self.store_layer_file(input_layer_2, file_path)
+        tmp_dir = tempfile.mkdtemp()
+        file_path = input_layer_2.download_to_working_directory(tmp_dir)
+        self.assertTrue(os.path.exists(file_path))
+        self.assertTrue(file_path.startswith(
+            os.path.join(
+                tmp_dir,
+                input_layer_2.get_component_type_display(),
+                os.path.basename(
+                    input_layer_2.file.name.replace('.zip', '_zip')
+                ),
+                'shops_poly'
+            )
+        ))
+
+        # Test invalid
+        input_layer_3 = InputLayerF.create(
+            owner=owner,
+            layer_type=InputLayer.LayerTypes.VECTOR,
+            component_type=InputLayer.ComponentTypes.MASK_LAYER
+        )
+        file_path = absolute_path(
+            'cplus_api', 'tests', 'data',
+            'mask_layers', 'shapefile_invalid.zip'
+        )
+        self.store_layer_file(input_layer_3, file_path)
+        tmp_dir = tempfile.mkdtemp()
+        file_path = input_layer_3.download_to_working_directory(tmp_dir)
+        self.assertIsNone(file_path)
+
+        # Test file not available
+        input_layer_4 = InputLayerF.create(
+            owner=owner,
+            layer_type=InputLayer.LayerTypes.VECTOR,
+            component_type=InputLayer.ComponentTypes.MASK_LAYER
+        )
+        file_path = absolute_path(
+            'cplus_api', 'tests', 'data',
+            'models', 'test_model_1.tif'
+        )
+        tmp_dir = tempfile.mkdtemp()
+        file_path = input_layer_4.download_to_working_directory(tmp_dir)
+        self.assertIsNone(file_path)

django_project/cplus_api/tests/test_utils.py

@@ -2,7 +2,11 @@
 import json
 import datetime
 from django.test import TestCase
-from cplus_api.utils.api_helper import todict, CustomJsonEncoder
+from cplus_api.utils.api_helper import (
+    todict,
+    CustomJsonEncoder,
+    get_layer_type
+)
 
 
 class SampleObj:
@@ -42,6 +46,19 @@ def test_list_dict(self):
         ]
         self.assertEqual(todict(test_obj, SampleObj), expected_value)
 
+    def test_get_layer_type(self):
+        filepath = '/tmp/test/file.tif'
+        self.assertEqual(get_layer_type(filepath), 0)
+
+        filepath = '/tmp/test/file.zip'
+        self.assertEqual(get_layer_type(filepath), 1)
+
+        filepath = '/tmp/test/file.shp'
+        self.assertEqual(get_layer_type(filepath), 1)
+
+        filepath = '/tmp/test/file.netcdf'
+        self.assertEqual(get_layer_type(filepath), -1)
+
 
 class TestCustomJSONEncoder(TestCase):
     def test_uuid(self):

django_project/cplus_api/utils/api_helper.py

@@ -260,3 +260,16 @@ def todict(obj, classkey=None):
         return data
     else:
         return obj
+
+
+def get_layer_type(file_path: str):
+    """
+    Get layer type code from file path
+    """
+    file_name, file_extension = os.path.splitext(file_path)
+    if file_extension.lower() in ['.tif', '.tiff']:
+        return 0
+    elif file_extension.lower() in ['.geojson', '.zip', '.shp']:
+        return 1
+    else:
+        return -1

django_project/cplus_api/utils/worker_analysis.py

@@ -19,8 +19,13 @@
 from cplus.tasks.analysis import ScenarioAnalysisTask
 from cplus.utils.conf import Settings
 from cplus_api.models.scenario import ScenarioTask
-from cplus_api.models.layer import BaseLayer, OutputLayer, InputLayer
-from cplus_api.utils.api_helper import convert_size, todict, CustomJsonEncoder
+from cplus_api.models.layer import OutputLayer, InputLayer
+from cplus_api.utils.api_helper import (
+    convert_size,
+    todict,
+    CustomJsonEncoder,
+    get_layer_type
+)
 from cplus_api.utils.default import DEFAULT_VALUES
 
 logger = logging.getLogger(__name__)
@@ -64,8 +69,7 @@ def __init__(self, scenario_name, scenario_desc, extent,
                  analysis_activities, priority_layers,
                  priority_layer_groups,
                  snapping_enabled=False, snap_layer_uuid='',
-                 pathway_suitability_index=
-                 DEFAULT_VALUES.pathway_suitability_index,
+                 pathway_suitability_index=DEFAULT_VALUES.pathway_suitability_index,  # noqa
                  carbon_coefficient=DEFAULT_VALUES.carbon_coefficient,
                  snap_rescale=DEFAULT_VALUES.snap_rescale,
                  snap_method=DEFAULT_VALUES.snap_method,
@@ -342,31 +346,37 @@ def create_and_upload_output_layer(
         is_final_output: bool, group: str,
         output_meta: dict = None) -> OutputLayer:
     filename = os.path.basename(file_path)
-    cog_name = (
-        f"{os.path.basename(file_path).split('.')[0]}"
-        f"_COG."
-        f"{os.path.basename(file_path).split('.')[1]}"
-    )
-    cog_path = os.path.join(
-        os.path.dirname(file_path),
-        cog_name
-    )
-    subprocess.run(
-        f"gdal_translate -of COG -co COMPRESS=DEFLATE {file_path} {cog_path}",
-        shell=True
-    )
+    if get_layer_type(file_path) == 0:
+        cog_name = (
+            f"{os.path.basename(file_path).split('.')[0]}"
+            f"_COG."
+            f"{os.path.basename(file_path).split('.')[1]}"
+        )
+        final_output_path = os.path.join(
+            os.path.dirname(file_path),
+            cog_name
+        )
+        subprocess.run(
+            (
+                f"gdal_translate -of COG -co COMPRESS=DEFLATE"
+                f" {file_path} {final_output_path}"
+            ),
+            shell=True
+        )
+    else:
+        final_output_path = file_path
     output_layer = OutputLayer.objects.create(
         name=filename,
         created_on=timezone.now(),
         owner=scenario_task.submitted_by,
-        layer_type=BaseLayer.LayerTypes.RASTER,
-        size=os.stat(cog_path).st_size,
+        layer_type=get_layer_type(file_path),
+        size=os.stat(final_output_path).st_size,
         is_final_output=is_final_output,
         scenario=scenario_task,
         group=group,
         output_meta={} if not output_meta else output_meta
     )
-    with open(cog_path, 'rb') as output_file:
+    with open(final_output_path, 'rb') as output_file:
         output_layer.file.save(filename, output_file)
     return output_layer