Gitlab Community Edition Instance

Commit 47e34cc2 authored by mhellka's avatar mhellka
Browse files

Merge branch 'feature-filedownload' into 'master'

Return FileDownload from CDStar.get_file

See merge request !6
parents 8456ba9b 00064e44
from pycdstar3.client import * # noqa: F401, F403
__version__ = "3.0.dev0"
from pycdstar3.client import * # noqa: F401, F403
......@@ -57,16 +57,15 @@ def get(ctx, args): # noqa: C901
close = out.close
try:
rs = client.get_file(vault, archive, file, offset=offset)
size = int(rs.headers["Content-Length"])
dl = client.get_file(vault, archive, file, offset=offset)
if ispipe and out.isatty() and not rs.headers["Content-Type"].startswith("text/"):
raise CliError("Not printing binary data ({}) a terminal".format(rs.headers["Content-Type"]))
if ispipe and out.isatty() and not dl.type.startswith("text/"):
raise CliError("Not printing binary data ({}) a terminal".format(dl.type))
if progress and not printer.quiet:
from tqdm import tqdm
pbar = tqdm(total=size+offset, initial=offset, unit='b', unit_scale=True, unit_divisor=1024, dynamic_ncols=True,
file=printer.file)
pbar = tqdm(total=dl.size+offset, initial=offset, unit='b', unit_scale=True, unit_divisor=1024,
dynamic_ncols=True, file=printer.file)
def write(chunk):
pbar.update(len(chunk))
......@@ -76,9 +75,9 @@ def get(ctx, args): # noqa: C901
pbar.close()
out.close()
printer.v("Downloading {!r} from /{}/{} ({})", file, vault, archive, hbytes(size + offset))
printer.v("Downloading {} ({})", file, hbytes(dl.size + offset))
for chunk in rs.iter_content(chunk_size=64*1024):
for chunk in dl.iter_content():
write(chunk)
close()
......@@ -88,7 +87,7 @@ def get(ctx, args): # noqa: C901
else:
os.rename(partfile, dst)
printer.v("Done!", file, vault, archive, hbytes(size + offset))
printer.v("Done!", file, vault, archive, hbytes(dl.size + offset))
except (KeyboardInterrupt, Exception):
if close:
......
......@@ -55,6 +55,7 @@ class CDStar:
rs = self._session.request(
method, self.url + '/'.join(path), **options)
if rs.ok or (expect_status and rs.status_code in expect_status):
return rs
......@@ -110,7 +111,7 @@ class CDStar:
raise RuntimeError("No transaction running")
try:
self.rest("POST", "_tx", self._tx['id'])
self.raw("POST", "_tx", self._tx['id'])
self._tx = None
except Exception:
self.rollback()
......@@ -118,8 +119,10 @@ class CDStar:
def rollback(self) -> None:
""" Rollback the current transaction, if any. Do nothing otherwise. """
if self._tx:
self.rest("DELETE", "_tx", self._tx['id'])
try:
if self._tx:
self.raw("DELETE", "_tx", self._tx['id'])
finally:
self._tx = None
def keepalive(self):
......@@ -208,13 +211,12 @@ class CDStar:
return self.rest("PUT", vault, archive, _fix_filename(name), data=source,
headers={'Content-Type': type or "application/x-autodetect"})
def get_file(self, vault, archive, name, offset=0):
def get_file(self, vault, archive, name, offset=0) -> "FileDownload":
""" Return a stream-able response object representing the requested file. """
# TODO: Return a wrapper class with convenience properties for size,
# save_as, __iter__
headers = {'Range': "bytes={}-".format(offset)} if offset > 0 else {}
return self.raw("GET", vault, archive, _fix_filename(name), stream=True, headers=headers)
name = _fix_filename(name)
rs = self.raw("GET", vault, archive, name, stream=True, headers=headers)
return FileDownload(vault, archive, name, rs)
def search(self, vault, q, order=None, limit=0, scroll=None, groups=None):
""" Perform a search and return the SearchResults document.
......@@ -246,6 +248,68 @@ class CDStar:
break
class FileDownload:
""" Wrapper for streamed file downloads.
The file content can only be read once.
"""
def __init__(self, vault, archive, name, rs: requests.Response):
self.vault = vault
self.archive = archive
self.name = name
self.rs = rs
self.read = rs.raw.read
@property
def basename(self):
return os.path.basename(self.name)
@property
def type(self):
return self.rs.headers["Content-Type"]
@property
def is_partial(self):
return self.rs.status_code == 206
@property
def size(self):
return int(self.rs.headers["Content-Length"])
def __iter__(self):
""" Iterate over chunks of data (NOT lines) """
return self.iter_content()
def iter_content(self, buffsize=1024 * 64):
""" Iterate over chunks of data (NOT lines) """
return self.rs.iter_content(chunk_size=buffsize)
def save_as(self, target):
""" Save this download to a file (str, Path or file-like)"""
if not hasattr(target, 'write'):
with open(target, 'wb') as fp:
return self.save_as(fp)
for chunk in self.iter_content():
target.write(chunk)
def close(self):
self.rs.close()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
self.close()
def __repr__(self):
return "FileDownload({}/{} name={!r} size={})".format(self.vault, self.archive, self.name, self.size)
def readall(self):
""" Read the entire download into memory and return a single large byte object. """
return self.rs.content
class ApiError(Exception):
def __init__(self, json):
self.json = json
......@@ -269,9 +333,6 @@ class ApiError(Exception):
def __repr__(self):
return '{0.error}({0.status}, msg="{0.message}")'.format(self)
def __str__(self):
return '{0.error}({0.status}, msg="{0.message}")'.format(self)
def pretty(self):
err = "API Error: {} ({})\n".format(self.error, self.status)
err += "Message: {}\n".format(self.message)
......
......@@ -67,8 +67,8 @@ def test_crud_file(): # pragma: no cover
fget = c.get_file(VAULT, a['id'], "test.py")
fp.seek(0)
assert fp.read() == fget.content
assert "text/x-python" == fget.headers['Content-Type']
assert fp.read() == fget.readall()
assert "text/x-python" == fget.type
# no commit
# Mockserver tests:
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment