def unpack_and_create_local_message(msg, local_dir, unpack=None, delete=False):
def unpack_callback(var):
if not var['uid'].endswith(unpack):
return var
dirname, filename = os.path.split(var.pop('uri'))
basename, ext = os.path.splitext(filename)
packname = var.pop('uid')
new_names = unpackers[unpack](os.path.join(local_dir, packname),
delete)
var['dataset'] = [
dict(uid=nn, uri=os.path.join(local_dir, nn)) for nn in new_names
]
return var
if unpack is not None:
lmsg_data = translate_dict(msg.data, ('uri', 'uid'), unpack_callback)
if 'dataset' in lmsg_data:
lmsg_type = 'dataset'
elif 'collection' in lmsg_data:
lmsg_type = 'collection'
else:
lmsg_type = 'file'
else:
lmsg_data = msg.data.copy()
lmsg_type = msg.type
return Message(msg.subject, lmsg_type, data=lmsg_data)
def test_item_translator(self):
"""Test translating dict items."""
orig = (
'pytroll://tm1 file [email protected] 2018-10-25T01:15:54.752065 v1.01 application/json '
'{"sensor": "viirs", "format": "SDR", "variant": "DR", "uid": "bla.tar", "uri": "/home/user/bla.tar"}'
)
dest_dir = '/tmp'
expected = (
'pytroll://tm1 file [email protected] 2018-10-25T01:15:54.752065 v1.01 application/json '
'{"sensor": "viirs", "format": "SDR", "variant": "DR", "uid": "bla.png", "uri": "/tmp/bla.png"}'
)
def dummy_cb(var, k):
dirname, filename = os.path.split(var[k])
basename, ext = os.path.splitext(filename)
if dirname:
dest = os.path.join(dest_dir, basename + '.png')
else:
dest = basename + '.png'
var[k] = dest
return var
expected_dict = Message(rawstr=expected).data
res = translate_dict_item(Message(rawstr=orig).data, 'uri', dummy_cb)
res = translate_dict_item(res, 'uid', dummy_cb)
self.assertDictEqual(expected_dict, res)
expected = (
'pytroll://tm1 file [email protected] 2018-10-25T01:15:54.752065 v1.01 application/json '
'{"sensor": "viirs", "format": "SDR", "variant": "DR", "dataset": [{"uid": "bla1.png", "uri": "/tmp/bla1.png"}, {"uid": "bla2.png", "uri": "/tmp/bla2.png"}]}'
) # noqa
def dummy_cb(var):
if not var['uid'].endswith('.tar'):
return var
dirname, filename = os.path.split(var.pop('uri'))
basename, ext = os.path.splitext(filename)
new_names = [basename + str(i) + '.png' for i in range(1, 3)]
var.pop('uid')
var['dataset'] = [
dict(uid=nn, uri=os.path.join(dest_dir, nn))
for nn in new_names
]
return var
expected_dict = Message(rawstr=expected).data
res = translate_dict(
Message(rawstr=orig).data, ('uri', 'uid'), dummy_cb)
self.assertDictEqual(expected_dict, res)
def test_translate_dict_item_to_dataset(self):
"""Test translating dict items to dataset."""
orig = (
'pytroll://tm1 file [email protected] 2018-10-25T01:15:54.752065 v1.01 application/json '
'{"sensor": "viirs", "format": "SDR", "variant": "DR", "uid": "bla.tar", "uri": "/home/user/bla.tar"}'
)
expected = (
'pytroll://tm1 file [email protected] 2018-10-25T01:15:54.752065 v1.01 application/json '
'{"sensor": "viirs", "format": "SDR", "variant": "DR", "dataset": [{"uid": "bla1.png", "uri": "/tmp/bla1.png"}, {"uid": "bla2.png", "uri": "/tmp/bla2.png"}]}'
) # noqa
expected_dict = Message(rawstr=expected).data
res = translate_dict(
Message(rawstr=orig).data, ('uri', 'uid'),
dummy_callback_translate_dict_item_to_dataset)
self.assertDictEqual(expected_dict, res)
def make_uris(msg, destination, login=None):
duri = urlparse(destination)
scheme = duri.scheme or 'ssh'
dest_hostname = duri.hostname or socket.gethostname()
if socket.gethostbyname(dest_hostname) in get_local_ips():
scheme_, host_ = "ssh", dest_hostname # local file
else:
scheme_, host_ = scheme, dest_hostname # remote file
if login:
# Add (only) user to uri.
host_ = login.split(":")[0] + "@" + host_
def uri_callback(var):
uid = var['uid']
path = os.path.join(duri.path, uid)
var['uri'] = urlunparse((scheme_, host_, path, "", "", ""))
return var
msg.data = translate_dict(msg.data, ('uri', 'uid'), uri_callback)
return msg
def unpack_and_create_local_message(msg, local_dir, **kwargs):
"""Unpack the file(s) given in the message, and return an updated message."""
def unpack_callback(var, **kwargs):
unpack = kwargs.get('compression')
endings = COMPRESSED_ENDINGS[unpack]
is_compressed = any(
[var['uid'].endswith(ending) for ending in endings])
if not is_compressed:
return var
packname = var.pop('uid')
del var['uri']
new_names = unpackers[unpack](os.path.join(local_dir, packname),
**kwargs)
if kwargs.get("delete"):
LOGGER.debug("Deleting %s", os.path.join(local_dir, packname))
os.remove(os.path.join(local_dir, packname))
if isinstance(new_names, tuple):
var['dataset'] = [
dict(uid=os.path.basename(nn), uri=os.path.join(local_dir, nn))
for nn in new_names
]
else:
var['uid'] = os.path.basename(new_names)
var['uri'] = os.path.join(local_dir, new_names)
return var
if kwargs.get('compression') in COMPRESSED_ENDINGS:
lmsg_data = translate_dict(msg.data, ('uri', 'uid'), unpack_callback,
**kwargs)
if 'dataset' in lmsg_data:
lmsg_type = 'dataset'
elif 'collection' in lmsg_data:
lmsg_type = 'collection'
else:
lmsg_type = 'file'
else:
lmsg_data = msg.data.copy()
lmsg_type = msg.type
return Message(msg.subject, lmsg_type, data=lmsg_data)