Commit aa374d5b authored by Derek Homeier's avatar Derek Homeier
Browse files

Unpack obsrec description into Quantity QueryResponseColumns

parent 00357eb4
Pipeline #2757 failed with stage
in 4 minutes and 5 seconds
......@@ -261,6 +261,17 @@ class KISClient(BaseClient):
results : `~sunpy.net.base_client.QueryResponseTable`
A `QueryResponseTable` instance containing the query result.
"""
# Unit converters for 'description' fields
converters = {'ATMOS': lambda x: x * u.cm,
'DATE_': lambda x: Time(x['$date'] * 1e-3, format='unix'),
'EXPTI': lambda x: x * u.s,
'HPLN_': lambda x: x * u.arcsec,
'HPLT_': lambda x: x * u.arcsec,
'S_RES': lambda x: x * u.arcsec,
'THETA': lambda x: x * u.deg,
'T_RES': lambda x: x * u.deg,
'WAVE_': lambda x: x * u.nm}
query = and_(*query)
# Not very extensively tested on hierarchical queries
......@@ -277,6 +288,10 @@ class KISClient(BaseClient):
obs_dict = json.loads(response.read()).get('_embedded', [])
if len(obs_dict) > 0 and 'links' in obs_dict[0]:
l1_data = obs_dict[0]['links'].get('l1_data', [])
obs_dict[0].update(obs_dict[0].pop('description', dict()))
for k in obs_dict[0].keys():
if k[:5] in converters:
obs_dict[0][k] = converters[k[:5]](obs_dict[0][k])
obs_rec = []
for exp in l1_data:
obs_rec.append(copy.deepcopy(obs_dict[0]))
......@@ -316,11 +331,11 @@ class KISClient(BaseClient):
ext = 'json'
for row in query_results:
inst = row['description']['INSTRUMENT']
inst = row['INSTRUMENT']
oid = row['links']['l1_data']['$oid']
# Content-Disposition header default is "{row['_id']['$oid']}/{oid}.{ext}" (no '.json').
# rowpath = row['_id']['$oid']
filepath = os.path.join(row['description']['OBS_NAME'], f"{oid}.{ext}")
filepath = os.path.join(row['OBS_NAME'], f"{oid}.{ext}")
url = f"{self._BASE_URL}{inst}_l1_data.files/{oid}{binfile}"
downloader.enqueue_file(url, filename=str(path).format(file=filepath,
**row.response_block_map))
......
......@@ -8,6 +8,7 @@ import json
import parfive
import astropy.units as u
from astropy import table
from astropy.io import fits
from astropy.time import Time
import sunpy
......@@ -69,13 +70,11 @@ def test_docker(client):
res = client.search(a.Instrument("GRIS") & a.sdc.ObsName('gris_20140426_000'))
assert isinstance(res, QueryResponseTable)
assert len(res) == 105
description = res[0].get('description')
assert len(description) == 34
assert description['INSTRUMENT'] == 'gris'
assert description['TELESCOPE'] == 'GREGOR'
assert description['BTYPE'] == 'phot.count'
assert description['DATE_BEG']['$date'] == 1398505619000
assert description['DATE_END']['$date'] == 1398506021300
assert res[0]['INSTRUMENT'] == 'gris'
assert res[0]['TELESCOPE'] == 'GREGOR'
assert res[0]['BTYPE'] == 'phot.count'
assert res[0]['DATE_BEG'].unix == 1398505619.000
assert res[0]['DATE_END'].unix == 1398506021.300
file_ids = [ld['links']['l1_data']['$oid'] for ld in res]
assert len(file_ids) == 105
......@@ -89,7 +88,7 @@ def test_docker(client):
date = a.Time("2014/04/26 01:00", "2014/04/26 22:00")
downloader = parfive.Downloader()
inst = res[0]['description']['INSTRUMENT']
inst = res[0]['INSTRUMENT']
rowpath = f"{res[0]['_id']['$oid']}"
binfile = ''
ext = 'json'
......@@ -166,7 +165,7 @@ def test_search(client):
res = client.search(query)
assert isinstance(res, QueryResponseTable)
assert len(res) == 1
assert 'description' in res.colnames
assert 'INSTRUMENT' in res.colnames
else:
with pytest.raises(URLError, match=rf"{_dockerexc('bbi')}"
r"{'description.DATE_BEG':{'.lte':{'.date':'2017-05-22T22:00:00.000'}}},"
......@@ -180,7 +179,7 @@ def test_search(client):
if HAS_DOCKERTEST:
res = client.search(query)
assert len(res) == 1
assert 'description' in res.colnames
assert 'INSTRUMENT' in res.colnames
else:
with pytest.raises(URLError, match=rf"{_dockerexc('lars')}"
r"{'description.HPLT_TAN_MIN':{'.lte':12}},"
......@@ -192,7 +191,7 @@ def test_search(client):
if HAS_DOCKERTEST:
res = client.search(*query)
assert len(res) == 1
assert 'description' in res.colnames
assert 'INSTRUMENT' in res.colnames
else:
with pytest.raises(URLError, match=rf"{_dockerexc('lars')}"
r"{'description.HPLT_TAN_MIN':{'.lte':12}},"
......@@ -204,11 +203,9 @@ def test_search(client):
if HAS_DOCKERTEST:
res = client.search(query)
assert len(res) == 362
assert 'description' in res.colnames
assert 'THETA' in res[0]['description']
theta = [obs['description']['THETA'] for obs in res]
assert (min(theta[:100]) >= 50) & (max(theta[:100]) <= 85)
assert res[100]['description']['POL_STATES'] == 'IQUV'
assert 'THETA' in res.colnames
assert (min(res['THETA'][:100]) >= 50*u.deg) & (max(res['THETA'][:100]) <= 85*u.deg)
assert res[100]['POL_STATES'] == 'IQUV'
else:
# Will raise on first of multi-part OR queries; somehow switches INSTRUMENT and THETA.
with pytest.raises(URLError, match=rf"{_dockerexc('gris')[:80]}") as exc:
......@@ -220,10 +217,8 @@ def test_search(client):
if HAS_DOCKERTEST:
res = client.search(*query)
assert len(res) == 163
assert 'description' in res.colnames
assert 'THETA' in res[0]['description']
theta = [obs['description']['THETA'] for obs in res]
assert (min(theta) >= 50) & (max(theta) <= 85)
assert 'THETA' in res.colnames
assert (min(res['THETA']) >= 50*u.deg) & (max(res['THETA']) <= 85*u.deg)
else:
# Will raise on first of multi-part OR queries; somehow switches INSTRUMENT and THETA.
with pytest.raises(URLError, match=rf"{_dockerexc('LARS')[:80]}") as exc:
......@@ -238,37 +233,37 @@ def test_fido_search():
if HAS_DOCKERTEST:
res = Fido.search(a.Instrument("GRIS") & a.sdc.Theta(50*u.deg, 80*u.deg))
assert len(res['kis']) == 162
theta = [obs['description']['THETA'] for obs in res['kis']]
assert (min(theta) >= 50) & (max(theta) <= 80)
assert (min(res['kis']['THETA']) >= 50*u.deg) & (max(res['kis']['THETA']) <= 80*u.deg)
res = Fido.search(a.Instrument("GRIS"), a.sdc.Theta(50*u.deg, 80*u.deg))
assert len(res['kis']) == 162
theta = [obs['description']['THETA'] for obs in res['kis']]
assert (min(theta) >= 50) & (max(theta) <= 80)
assert (min(res['kis']['THETA']) >= 50*u.deg) & (max(res['kis']['THETA']) <= 80*u.deg)
date = a.Time("2017/05/12 01:40", "2017/05/16 19:00")
res = Fido.search(a.Instrument("LARS") & date)
assert len(res['kis']) > 0
date_beg = [obs['description']['DATE_BEG']['$date'] for obs in res['kis']]
date_end = [obs['description']['DATE_END']['$date'] for obs in res['kis']]
assert max(date_beg) < date.end.unix * 1000
assert min(date_end) > date.start.unix * 1000
assert max(res['kis']['DATE_BEG']) < date.end
assert min(res['kis']['DATE_END']) > date.start
res = Fido.search(two_inst, a.sdc.Theta(50*u.deg, 80*u.deg))
assert len(res['kis']) == 2
assert len(res['kis'][0]) == 1
assert len(res['kis'][1]) == 162
assert res['kis'][0][0]['description']['INSTRUMENT'] == 'lars'
theta = [obs['description']['THETA'] for obs in res['kis'][0]]
assert (min(theta) >= 50) & (max(theta) <= 80)
assert res['kis'][1, 0]['description']['INSTRUMENT'] == 'gris'
theta = [obs['description']['THETA'] for obs in res['kis'][1]]
assert (min(theta) >= 50) & (max(theta) <= 80)
assert res['kis'][0][0]['INSTRUMENT'] == 'lars'
assert res['kis'][1, 0]['INSTRUMENT'] == 'gris'
assert (min(res[0]['kis']['THETA']) >= 50*u.deg) & (max(res[0]['kis']['THETA']) <= 80*u.deg)
assert (min(res[1]['kis']['THETA']) >= 50*u.deg) & (max(res[1]['kis']['THETA']) <= 80*u.deg)
theta = [obs['THETA'] for obs in res['kis'][0]]
assert (min(theta) >= 50*u.deg) & (max(theta) <= 80*u.deg)
theta = [obs['THETA'] for obs in res['kis'][1]]
assert (min(theta) >= 50*u.deg) & (max(theta) <= 80*u.deg)
date = a.Time("2016/08/26 16:25", "2016/08/26 16:45")
res = Fido.search(a.Instrument("GRIS"), a.sdc.PolStates('iquv'), date)
assert len(res['kis']) == 400
assert res['kis'][0]['description']['POL_STATES'] == 'IQUV'
assert all(res['kis'][0]['POL_STATES'] == 'IQUV')
assert max(res['kis']['DATE_BEG']) < date.end
assert min(res['kis']['DATE_END']) > date.start
else:
with pytest.raises(URLError, match=rf"{_dockerexc('gris')}"
r"{'description.THETA':{'.gte':50,'.lte':80}}"):
......@@ -292,25 +287,23 @@ def test_fido_fetch():
date = a.Time("2017/05/22 08:45", "2017/05/22 08:55")
res = Fido.search(a.Instrument("BBI"), date)
assert len(res['kis']) == 10
desc = [r['description'] for r in res['kis']]
assert max([d['DATE_BEG']['$date'] for d in desc]) < date.end.unix * 1000
assert min([d['DATE_END']['$date'] for d in desc]) > date.start.unix * 1000
assert max(res['kis']['DATE_BEG']) < date.end
assert min(res['kis']['DATE_END']) > date.start
files = Fido.fetch(res['kis'])
assert len(files) == 10
for i, filepath in enumerate(files):
assert dirnames(filepath)[-1] == desc[::-1][i]['OBS_NAME']
assert dirnames(filepath)[-1] == res['kis'][::-1][i]['OBS_NAME']
meta = json.load(open(filepath))
assert meta['_id']['$oid'] == os.path.splitext(os.path.basename(filepath))[0]
assert date.start.isot[:12] in meta['metadata']['header']['DATE-BEG']
assert date.start < Time(meta['metadata']['header']['DATE-BEG']) < date.end
assert meta['metadata']['header']['FILENAME'].split('-')[0] in desc[::-1][i]['OBS_NAME']
assert meta['metadata']['header']['FILENAME'].split('-')[0] in res['kis'][9-i]['OBS_NAME']
files = Fido.fetch(res['kis'][:3], binary=True)
assert len(files) == 3
for i, filepath in enumerate(files):
assert dirnames(filepath)[-1] == desc[::-1][i]['OBS_NAME']
assert dirnames(filepath)[-1] == res['kis'][::-1][i]['OBS_NAME']
hdulist = fits.open(filepath)
assert hdulist[0].header.get('TELESCOP') == 'GREGOR'
assert hdulist[0].header.get('INSTRUME') == 'BBI'
......@@ -335,9 +328,8 @@ def test_fido_fetch_2():
res = Fido.search((a.Instrument("GRIS") | a.Instrument("LARS")) & date)
assert len(res['kis']) == 2
assert len(res['kis'][0]) == 400
assert res['kis'][0][0]['description']['DATE_BEG']['$date'] < date.end.unix * 1000
assert res['kis'][0][0]['description']['DATE_END']['$date'] > date.start.unix * 1000
assert max(res['kis'][0]['DATE_BEG']) < date.end
assert min(res['kis'][0]['DATE_END']) > date.start
files = Fido.fetch(res['kis'][0, :100], binary=False)
assert len(files) == 100
......@@ -372,19 +364,22 @@ def test_fido_fetch_2():
assert len(res['kis']) == 2
assert len(res['kis'][0]) == 300
assert len(res['kis'][1]) == 1
desc = [r['description'] for r in res['kis'][0]] + [r['description'] for r in res['kis'][1]]
dirs = [d['OBS_NAME'] for d in desc]
assert len(desc) == 301
assert max([d['DATE_BEG']['$date'] for d in desc]) < date.end.unix * 1000
assert min([d['DATE_END']['$date'] for d in desc]) > date.start.unix * 1000
desc = table.vstack(res['kis'])
assert len(desc) == 301
assert max(res['kis'][0]['DATE_BEG']) < date.end
assert min(res['kis'][0]['DATE_END']) > date.start
assert max(res['kis'][1]['DATE_BEG']) < date.end
assert min(res['kis'][1]['DATE_END']) > date.start
assert max(desc['DATE_BEG']) < date.end
assert min(desc['DATE_END']) > date.start
files = Fido.fetch(res['kis'][:, :100], binary=False)
assert len(files.errors) == 0
assert len(files) == 101
assert files[0].endswith('.json')
for filepath in files:
assert dirnames(filepath)[-1] in dirs
assert dirnames(filepath)[-1] in desc['OBS_NAME']
meta = json.load(open(filepath))
assert meta['_id']['$oid'] == os.path.splitext(os.path.basename(filepath))[0]
assert date.start.iso[:10] in meta['metadata']['header']['DATE-OBS']
......@@ -440,7 +435,7 @@ def test_all_queries(client, query):
if HAS_DOCKERTEST:
res = client.search(a.Instrument("GRIS") & query)
if len(res) > 0:
assert 'description' in res.colnames
assert res[0]['INSTRUMENT'] == 'gris'
else:
with pytest.raises(URLError, match=rf"{_dockerexc('gris')}"
rf"{{'description.*{query.type_name.upper()}"):
......@@ -456,10 +451,8 @@ def test_range(client):
assert client._can_handle_query(query)
if HAS_DOCKERTEST:
res = client.search(query)
wave_min = [obs['WAVELENGTH_MIN'] for obs in res['description']]
wave_max = [obs['WAVELENGTH_MAX'] for obs in res['description']]
assert max(wave_min) <= 1250
assert min(wave_max) >= 1080
assert max(res['WAVELENGTH_MIN']) <= 1250 * u.nm
assert min(res['WAVELENGTH_MAX']) >= 1080 * u.nm
else:
with pytest.raises(URLError, match=rf"{_dockerexc('gris')}"
r"{'description.WAVELENGTH_MIN':{'.lte':1250}},"
......@@ -479,7 +472,7 @@ def test_full_range(client):
if HAS_DOCKERTEST:
res = client.search(query)
assert len(res) == 105
assert 'description' in res.colnames
assert 'INSTRUMENT' in res.colnames
else:
with pytest.raises(URLError, match=rf"{_dockerexc('GRIS')}"
r"{'description.DATE_BEG':{'.lte':{'.date':'2014-04-26T09:50:00.000'}}},"
......@@ -491,10 +484,8 @@ def test_full_range(client):
query = a.Instrument("LARS") & hplt
if HAS_DOCKERTEST:
res = client.search(query)
hplt_tan_min = [obs['HPLT_TAN_MIN'] for obs in res['description']]
hplt_tan_max = [obs['HPLT_TAN_MAX'] for obs in res['description']]
assert max(hplt_tan_min) <= 6
assert min(hplt_tan_max) >= -2
assert max(res['HPLT_TAN_MIN']) <= 6 * u.arcsec
assert min(res['HPLT_TAN_MAX']) >= -2 * u.arcsec
else:
with pytest.raises(URLError, match=rf"{_dockerexc('LARS')}"
r"{'description.HPLT_TAN_MIN':{'.lte':6}},"
......@@ -506,10 +497,8 @@ def test_full_range(client):
assert client._can_handle_query(query)
if HAS_DOCKERTEST:
res = client.search(query)
hplt_tan_min = [obs['HPLT_TAN_MIN'] for obs in res['description']]
hplt_tan_max = [obs['HPLT_TAN_MAX'] for obs in res['description']]
assert max(hplt_tan_min) <= -2
assert min(hplt_tan_max) >= 6
assert max(res['HPLT_TAN_MIN']) <= -2 * u.arcsec
assert min(res['HPLT_TAN_MAX']) >= 6 * u.arcsec
else:
with pytest.raises(URLError, match=rf"{_dockerexc('LARS')}"
r"{'description.HPLT_TAN_MIN':{'.lte':-2}},"
......
Markdown is supported
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