Skip to content

Commit 946e6c7

Browse files
committed
added a getAsZip and putAsZip function for both directories and single files; with tests
1 parent 2915fc8 commit 946e6c7

File tree

5 files changed

+51
-1
lines changed

5 files changed

+51
-1
lines changed

Algorithmia/datafile.py

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
from datetime import datetime
88
import os.path
99
import pkgutil
10+
import shutil
11+
import zipfile
1012

1113
from Algorithmia.util import getParentAndBase
1214
from Algorithmia.data import DataObject, DataObjectType
@@ -50,6 +52,18 @@ def getFile(self, as_path=False):
5052
else:
5153
return open(f.name)
5254

55+
def getAsZip(self):
56+
local_file_path = self.getFile(as_path=True)
57+
directory_path = tempfile.mkdtemp()
58+
with zipfile.ZipFile(local_file_path, 'r') as ziph:
59+
ziph.extractall(directory_path)
60+
if len(ziph.namelist()) > 1:
61+
output_path = directory_path
62+
else:
63+
filename = ziph.namelist()[0]
64+
output_path = os.path.join(directory_path, filename)
65+
return output_path
66+
5367
def getName(self):
5468
_, name = getParentAndBase(self.path)
5569
return name
@@ -145,6 +159,20 @@ def putNumpy(self, array):
145159
else:
146160
raise DataApiError("Attempted to .putNumpy() a file without numpy available, please install numpy.")
147161

162+
def putAsZip(self, path):
163+
temp = tempfile.NamedTemporaryFile(delete=False).name
164+
if os.path.isdir(path):
165+
with zipfile.ZipFile(temp, 'w') as ziph:
166+
for root, dirs, files in os.walk(path):
167+
for file in files:
168+
f_path = os.path.join(root, file)
169+
arc_path = os.path.relpath(os.path.join(root, file), path)
170+
ziph.write(f_path, arc_path)
171+
else:
172+
with zipfile.ZipFile(temp, 'w') as ziph:
173+
ziph.write(path)
174+
return self.putFile(temp)
175+
148176
def delete(self):
149177
# Delete from data api
150178
result = self.client.deleteHelper(self.url)
@@ -256,7 +284,7 @@ def __del__(self):
256284
filepath = self.local_file.name
257285
self.local_file.close()
258286
if self.cleanup:
259-
os.remove(filepath)
287+
os.remove(filepath)
260288

261289
def readable(self):
262290
return True

Test/datafile_test.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import Algorithmia
1010
import json
1111
from Algorithmia.datafile import DataFile, LocalDataFile, AdvancedDataFile
12+
from pathlib import Path
1213

1314
class DataFileTest(unittest.TestCase):
1415
def setUp(self):
@@ -151,5 +152,21 @@ def test_putJson_getJson(self):
151152
self.assertDictEqual(result, payload)
152153
self.assertEqual(str(result), str(payload))
153154

155+
def test_putZipDir_getZipDir(self):
156+
local_directory = os.path.join(os.getcwd(), "resources/zip_directory")
157+
remote_directory = "data://.my/empty/datafile.zip"
158+
df = AdvancedDataFile(self.client, remote_directory, cleanup=True)
159+
response = df.putAsZip(local_directory)
160+
self.assertEqual(response, df)
161+
162+
unzipped_local_path = df.getAsZip()
163+
self.assertTrue(os.path.isdir(unzipped_local_path))
164+
found_files = []
165+
for _, _, files in os.walk(unzipped_local_path):
166+
for file in files:
167+
found_files.append(file)
168+
self.assertEqual(len(found_files), 3)
169+
170+
154171
if __name__ == '__main__':
155172
unittest.main()
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
{"location": "root"}
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from .build_wait import get_build
2+
from .publish_algo import publish_algo
3+
from .test_algo import test_algo
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
{"foo": "bar"}

0 commit comments

Comments
 (0)