Merge pull request #626 from JustAnotherArchivist/fix-upload-size-hint

Fix automatic size hint on uploads

internetarchive/item.py

@@ -57,7 +57,7 @@
     iter_directory,
     json,
     norm_filepath,
-    recursive_file_count,
+    recursive_file_count_and_size,
     validate_s3_identifier,
 )
 
@@ -1194,11 +1194,13 @@ def upload(self, files,
 
         responses = []
         file_index = 0
-        if queue_derive and total_files is None:
-            if checksum:
-                total_files = recursive_file_count(files, item=self, checksum=True)
-            else:
-                total_files = recursive_file_count(files, item=self, checksum=False)
+        headers = headers or {}
+        if (queue_derive or not headers.get('x-archive-size-hint')) and total_files == 0:
+            total_files, total_size = recursive_file_count_and_size(files,
+                                                                    item=self,
+                                                                    checksum=checksum)
+            if not headers.get('x-archive-size-hint'):
+                headers['x-archive-size-hint'] = str(total_size)
         file_metadata = None
         for f in files:
             if isinstance(f, dict):

internetarchive/utils.py

@@ -216,15 +216,19 @@ def _get_tag_text(tag_name, xml_obj):
 
 
 def get_file_size(file_obj) -> int | None:
-    try:
-        file_obj.seek(0, os.SEEK_END)
-        size = file_obj.tell()
-        # Avoid OverflowError.
-        if size > sys.maxsize:
+    if is_filelike_obj(file_obj):
+        try:
+            file_obj.seek(0, os.SEEK_END)
+            size = file_obj.tell()
+            # Avoid OverflowError.
+            if size > sys.maxsize:
+                size = None
+            file_obj.seek(0, os.SEEK_SET)
+        except OSError:
             size = None
-        file_obj.seek(0, os.SEEK_SET)
-    except OSError:
-        size = None
+    else:
+        st = os.stat(file_obj)
+        size = st.st_size
     return size
 
 
@@ -237,11 +241,14 @@ def iter_directory(directory: str):
             yield (filepath, key)
 
 
-def recursive_file_count(files, item=None, checksum=False):
-    """Given a filepath or list of filepaths, return the total number of files."""
+def recursive_file_count_and_size(files, item=None, checksum=False):
+    """Given a filepath or list of filepaths, return the total number and size of files.
+    If `checksum` is `True`, skip over files whose MD5 hash matches any file in the `item`.
+    """
     if not isinstance(files, (list, set)):
         files = [files]
     total_files = 0
+    total_size = 0
     if checksum is True:
         md5s = [f.get('md5') for f in item.files]
     else:
@@ -264,24 +271,27 @@ def recursive_file_count(files, item=None, checksum=False):
             except (AttributeError, TypeError):
                 is_dir = False
         if is_dir:
-            for x, _ in iter_directory(f):
-                if checksum is True:
-                    with open(x, 'rb') as fh:
-                        lmd5 = get_md5(fh)
-                    if lmd5 in md5s:
-                        continue
-                total_files += 1
+            it = iter_directory(f)
         else:
+            it = [(f, None)]
+        for x, _ in it:
             if checksum is True:
                 try:
-                    with open(f, 'rb') as fh:
+                    with open(x, 'rb') as fh:
                         lmd5 = get_md5(fh)
                 except TypeError:
                     # Support file-like objects.
-                    lmd5 = get_md5(f)
+                    lmd5 = get_md5(x)
                 if lmd5 in md5s:
                     continue
+            total_size += get_file_size(x)
             total_files += 1
+    return total_files, total_size
+
+
+def recursive_file_count(*args, **kwargs):
+    """Like `recursive_file_count_and_size`, but returns only the file count."""
+    total_files, _ = recursive_file_count_and_size(*args, **kwargs)
     return total_files
 
 
@@ -294,6 +304,16 @@ def is_dir(obj) -> bool:
         return False
 
 
+def is_filelike_obj(obj) -> bool:
+    """Distinguish file-like from path-like objects"""
+    try:
+        os.fspath(obj)
+    except TypeError:
+        return True
+    else:
+        return False
+
+
 def reraise_modify(
     caught_exc: Exception,
     append_msg: str,

tests/cli/test_ia_upload.py

@@ -124,6 +124,27 @@ def test_ia_upload_size_hint(capsys, tmpdir_ch, nasa_mocker):
     assert 'Accept-Encoding:gzip, deflate' in err
 
 
+def test_ia_upload_automatic_size_hint_files(capsys, tmpdir_ch, nasa_mocker):
+    with open('foo', 'w') as fh:
+        fh.write('foo')
+    with open('bar', 'w') as fh:
+        fh.write('bar')
+
+    ia_call(['ia', 'upload', '--debug', 'nasa', 'foo', 'bar'])
+    out, err = capsys.readouterr()
+    assert 'x-archive-size-hint:6' in err
+
+def test_ia_upload_automatic_size_hint_dir(capsys, tmpdir_ch, nasa_mocker):
+    with open('foo', 'w') as fh:
+        fh.write('foo')
+    with open('bar', 'w') as fh:
+        fh.write('bar')
+
+    ia_call(['ia', 'upload', '--debug', 'nasa', '.'])
+    out, err = capsys.readouterr()
+    assert 'x-archive-size-hint:6' in err
+
+
 def test_ia_upload_unicode(tmpdir_ch, caplog):
     with open('தமிழ் - baz ∆.txt', 'w') as fh:
         fh.write('unicode foo')

tests/test_item.py

@@ -580,6 +580,40 @@ def test_upload_checksum(tmpdir, nasa_item):
             assert r.status_code is None
 
 
+def test_upload_automatic_size_hint(tmpdir, nasa_item):
+    with IaRequestsMock(assert_all_requests_are_fired=False) as rsps:
+        _expected_headers = deepcopy(EXPECTED_S3_HEADERS)
+        del _expected_headers['x-archive-size-hint']
+        _expected_headers['x-archive-size-hint'] = '15'
+        rsps.add(responses.PUT, S3_URL_RE,
+                 adding_headers=_expected_headers)
+
+        files = []
+        with open(os.path.join(tmpdir, 'file'), 'w') as fh:
+            fh.write('a')
+        files.append(os.path.join(tmpdir, 'file'))
+
+        os.mkdir(os.path.join(tmpdir, 'dir'))
+        with open(os.path.join(tmpdir, 'dir', 'file0'), 'w') as fh:
+            fh.write('bb')
+        with open(os.path.join(tmpdir, 'dir', 'file1'), 'w') as fh:
+            fh.write('cccc')
+        files.append(os.path.join(tmpdir, 'dir'))
+
+        with open(os.path.join(tmpdir, 'obj'), 'wb') as fh:
+            fh.write(b'dddddddd')
+            fh.seek(0, os.SEEK_SET)
+            files.append(fh)
+
+            _responses = nasa_item.upload(files,
+                                          access_key='a',
+                                          secret_key='b')
+        for r in _responses:
+            headers = {k.lower(): str(v) for k, v in r.headers.items()}
+            del headers['content-type']
+            assert headers == _expected_headers
+
+
 def test_modify_metadata(nasa_item, nasa_metadata):
     with IaRequestsMock(assert_all_requests_are_fired=False) as rsps:
         rsps.add(responses.POST, f'{PROTOCOL}//archive.org/metadata/nasa')