Fix up chunked uploading

siaskynet/__init__.py

@@ -24,7 +24,8 @@ class SkynetClient():
         get_skykeys
     )
     from ._upload import (
-        upload_file, upload_file_request, upload_file_request_with_chunks,
+        upload, upload_request,
+        upload_file, upload_file_request,
         upload_directory, upload_directory_request
     )
     # pylint: enable=import-outside-toplevel

siaskynet/_upload.py

@@ -18,6 +18,67 @@ def default_upload_options():
     return obj
 
 
+def upload(self, upload_data, custom_opts=None):
+    """Uploads the given generic data and returns the skylink."""
+
+    response = self.upload_request(upload_data, custom_opts)
+    sia_url = utils.uri_skynet_prefix() + response.json()["skylink"]
+    response.close()
+    return sia_url
+
+
+def upload_request(self, upload_data, custom_opts=None):
+    """Uploads the given generic data and returns the response object."""
+
+    opts = default_upload_options()
+    opts.update(self.custom_opts)
+    if custom_opts is not None:
+        opts.update(custom_opts)
+
+    if len(upload_data) == 1:
+        fieldname = opts['portal_file_fieldname']
+        # this appears to be missing in go sdk; maybe the server ignores it?
+        filename = opts['custom_filename']
+    else:
+        if not opts['custom_dirname']:
+            raise ValueError("custom_dirname must be set when "
+                             "uploading multiple files")
+        fieldname = opts['portal_directory_file_fieldname']
+        filename = opts['custom_dirname']
+
+    params = {
+        'filename': filename,
+        # 'skykeyname': opts['skykey_name'],
+        # 'skykeyid': opts['skyket_id'],
+    }
+
+    ftuples = []
+    for filename, data in upload_data.items():
+        ftuples.append((fieldname,
+                        (filename, data)))
+
+    if len(upload_data) == 1:
+        data = ftuples[0][1][1]
+        if hasattr(data, '__iter__') and (
+                not isinstance(data, bytes) and
+                not isinstance(data, str) and
+                not hasattr(data, 'read')):
+            # an iterator for chunked uploading
+            return self.execute_request(
+                "POST",
+                opts,
+                data=data,
+                headers={'Content-Type': 'application/octet-stream'},
+                params=params
+            )
+    return self.execute_request(
+        "POST",
+        opts,
+        files=ftuples,
+        params=params
+    )
+
+
 def upload_file(self, path, custom_opts=None):
     """Uploads file at path with the given options."""
 
@@ -48,41 +109,12 @@ def upload_file_request(self, path, custom_opts=None):
         return None
 
     with open(path, 'rb') as file_h:
-        filename = opts['custom_filename'] if opts['custom_filename'] \
-            else os.path.basename(file_h.name)
-        files = {opts['portal_file_fieldname']: (filename, file_h)}
+        if not opts['custom_filename']:
+            opts['custom_filename'] = os.path.basename(file_h.name)
 
-        return self.execute_request(
-            "POST",
-            opts,
-            files=files,
-        )
+        upload_data = {opts['custom_filename']: file_h}
 
-
-def upload_file_request_with_chunks(self, path, custom_opts=None):
-    """Posts request to upload file with chunks."""
-
-    opts = default_upload_options()
-    opts.update(self.custom_opts)
-    if custom_opts is not None:
-        opts.update(custom_opts)
-
-    path = os.path.normpath(path)
-    if not os.path.isfile(path):
-        print("Given path is not a file")
-        return None
-
-    filename = opts['custom_filename'] if opts['custom_filename'] else path
-    params = {filename: filename}
-    headers = {'Content-Type': 'application/octet-stream'}
-
-    return self.execute_request(
-        "POST",
-        opts,
-        data=path,
-        headers=headers,
-        params=params,
-    )
+        return self.upload_request(upload_data, opts)
 
 
 def upload_directory(self, path, custom_opts=None):
@@ -107,21 +139,13 @@ def upload_directory_request(self, path, custom_opts=None):
         print("Given path is not a directory")
         return None
 
-    ftuples = []
+    upload_data = {}
     basepath = path if path == '/' else path + '/'
-    files = list(utils.walk_directory(path).keys())
-    for filepath in files:
+    for filepath in utils.walk_directory(path):
         assert filepath.startswith(basepath)
-        ftuples.append((opts['portal_directory_file_fieldname'],
-                        (filepath[len(basepath):], open(filepath, 'rb'))))
-
-    dirname = opts['custom_dirname'] if opts['custom_dirname'] else path
+        upload_data[filepath[len(basepath):]] = open(filepath, 'rb')
 
-    params = {"filename": dirname}
+    if not opts['custom_dirname']:
+        opts['custom_dirname'] = path
 
-    return self.execute_request(
-        "POST",
-        opts,
-        files=ftuples,
-        params=params,
-    )
+    return self.upload_request(upload_data, opts)

tests/test_integration_upload.py

@@ -246,3 +246,45 @@ def test_upload_directory_custom_dirname():
 filename="file0"') == -1
 
     assert len(responses.calls) == 1
+
+
+@responses.activate
+def test_upload_file_chunks():
+    """Test uploading a file with chunks."""
+
+    src_file = "./testdata/file1"
+
+    # upload a file
+
+    responses.add(
+        responses.POST,
+        'https://siasky.net/skynet/skyfile',
+        json={'skylink': SKYLINK},
+        status=200
+    )
+
+    print("Uploading file "+src_file)
+
+    def chunker(filename):
+        with open(filename, 'rb') as file:
+            while True:
+                data = file.read(3)
+                if not data:
+                    break
+                yield data
+    chunks = chunker(src_file)
+    sialink2 = client.upload({'file1': chunks})
+    if SIALINK != sialink2:
+        sys.exit("ERROR: expected returned sialink "+SIALINK +
+                 ", received "+sialink2)
+    print("File upload successful, sialink: " + sialink2)
+
+    headers = responses.calls[0].request.headers
+    assert headers["Transfer-Encoding"] == "chunked"
+    assert headers["User-Agent"] == "python-requests/2.24.0"
+    assert "Authorization" not in headers
+
+    body = responses.calls[0].request.body
+    assert body is chunks
+
+    assert len(responses.calls) == 1