diff --git a/CHANGES.rst b/CHANGES.rst index c00ea684b..62e840407 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -6,6 +6,11 @@ * Using the six libray to address Python 2/3 compatibility issues +* AsyncTAPJob is now context aware + +* Improvement upload handling; it is no longer necessary to specifiy the type + of upload + 0.5.2 ---------------- Remove trailing ? from query urls diff --git a/pyvo/dal/query.py b/pyvo/dal/query.py index 3043a2b95..7f84c34db 100644 --- a/pyvo/dal/query.py +++ b/pyvo/dal/query.py @@ -36,6 +36,7 @@ import functools from astropy.extern import six +from astropy.table.table import Table if six.PY3: _mimetype_re = re.compile(b'^\w[\w\-]+/\w[\w\-]+(\+\w[\w\-]*)?(;[\w\-]+(\=[\w\-]+))*$') @@ -979,6 +980,143 @@ def __next__(self): next = __next__ +class Upload(object): + """ + This class represents a DALI Upload as described in + http://www.ivoa.net/documents/DALI/20161101/PR-DALI-1.1-20161101.html#tth_sEc3.4.5 + """ + + def __init__(self, name, content): + """ + Initialise the Upload object with the given parameters + + Parameters + ---------- + name : str + Tablename for use in queries + content : object + If its a file-like object, a string pointing to a local file, + a `DALResults` object or a astropy table, `is_inline` will be true + and it will expose a file-like object under `fileobj` + + Otherwise it exposes a URI under `uri` + """ + try: + self._is_file = os.path.isfile(content) + except Exception: + self._is_file = False + self._is_fileobj = hasattr(content, "read") + self._is_table = isinstance(content, Table) + self._is_resultset = isinstance(content, DALResults) + + self._inline = any(( + self._is_file, + self._is_fileobj, + self._is_table, + self._is_resultset, + )) + + self._name = name + self._content = content + + @property + def is_inline(self): + """ + True if the upload can be inlined + """ + return self._inline + + @property + def name(self): + return self._name + + def fileobj(self): + """ + A file-like object for a local resource + + Raises + ------ + ValueError + if theres no valid local resource + """ + + if not self.is_inline: + raise ValueError( + "Upload {name} doesn't refer to a local resource".format( + name = self.name)) + + # astropy table + if isinstance(self._content, Table): + from io import BytesIO + fileobj = BytesIO() + + self._content.write(output = fileobj, format = "votable") + fileobj.seek(0) + + return fileobj + elif isinstance(self._content, DALResults): + from io import BytesIO + fileobj = BytesIO() + + table = self._content.table + table.write(output = fileobj, format = "votable") + fileobj.seek(0) + + return fileobj + + fileobj = self._content + try: + fileobj = open(self._content) + finally: + return fileobj + + def uri(self): + """ + The URI pointing to the result + """ + + # TODO: use a async job base class instead of hasattr for inspection + if hasattr(self._content, "result_uri"): + self._content.raise_if_error() + uri = self._content.result_uri + else: + uri = six.text_type(self._content) + + return uri + + def query_part(self): + """ + The query part for use in DALI requests + """ + + if self.is_inline: + value = "{name},param:{name}" + else: + value = "{name},{uri}" + + return value.format(name = self.name, uri = self.uri()) + + +class UploadList(list): + """ + This class extends the native python list with utility functions for + upload handling + """ + + @classmethod + def fromdict(cls, dct): + """ + Constructs a upload list from a dictionary with table_name: content + """ + return cls(Upload(key, value) for key, value in dct.items()) + + def param(self): + """ + Returns a string suitable for use in UPLOAD parameters + """ + return ";".join(upload.query_part() for upload in self) + + if six.PY3: _image_mt_re = re.compile(b'^image/(\w+)') _text_mt_re = re.compile(b'^text/(\w+)') diff --git a/pyvo/dal/tap.py b/pyvo/dal/tap.py index ecd50accf..2085e19d2 100644 --- a/pyvo/dal/tap.py +++ b/pyvo/dal/tap.py @@ -12,8 +12,8 @@ from time import time, sleep from .query import ( - DALResults, DALQuery, DALService, DALServiceError, DALQueryError, - _votableparse) + DALResults, DALQuery, DALService, UploadList, + DALServiceError, DALQueryError, _votableparse) from ..tools import vosi, uws __all__ = ["search", "escape", @@ -123,8 +123,7 @@ def __init__(self, baseurl, query, mode="sync", language="ADQL", super(TAPQuery, self).__init__(baseurl, "TAP", "1.0") self._mode = mode if mode in ("sync", "async") else "sync" - self._uploads = uploads or {} - self._uploads = {k: _fix_upload(v) for k, v in self._uploads.items()} + self._uploads = UploadList.fromdict(uploads or {}) self["REQUEST"] = "doQuery" self["LANG"] = language @@ -135,13 +134,7 @@ def __init__(self, baseurl, query, mode="sync", language="ADQL", self["QUERY"] = query if self._uploads: - upload_param = ';'.join( - ['{0},{1}{2}'.format( - k, - 'param:' if v[0] == 'inline' else '', - v[1] if v[0] == 'uri' else k - ) for k, v in self._uploads.items()]) - self["UPLOAD"] = upload_param + self["UPLOAD"] = self._uploads.param() def getqueryurl(self): return '{0}/{1}'.format(self.baseurl, self._mode) @@ -173,8 +166,11 @@ def submit(self): """ url = self.getqueryurl() - files = {k: _fileobj(v[1]) for k, v in filter( - lambda x: x[1][0] == 'inline', self._uploads.items())} + files = { + upload.name: upload.fileobj() + for upload in self._uploads + if upload.is_inline + } r = requests.post(url, data = self, stream = True, files = files)