def query_collection(cid=None, cname=None, criteria=None):
from COMPS.Data import QueryCriteria
from COMPS.Data import AssetCollection
criteria = criteria or QueryCriteria().select_children('assets')
if cid:
return AssetCollection.get(id=cid, query_criteria=criteria)
criteria.where_tag('Name={}'.format(cname))
results = AssetCollection.get(query_criteria=criteria)
if len(results) >= 1: return results[0]
def get_asset_collection_by_id(collection_id, query_criteria=None):
query_criteria = query_criteria or QueryCriteria().select_children(
'assets')
try:
return AssetCollection.get(collection_id, query_criteria)
except (RuntimeError, ValueError):
return None
def get_asset_files_for_simulation_id(sim_id, paths, output_directory=None, flatten=False, remove_prefix=None):
"""
Obtains AssetManager-contained files from a given simulation.
:param sim_id: A simulation id to retrieve files from
:param file_paths: relative to the Assets folder
:param remove_prefix: if a prefix is given, will remove it from the paths
:param output_directory: Write requested files into this directory if specified
:param flatten: If true, all the files will be written to the root of output_directory. If false, dir structure will be kept
:return: Dictionary associating filename and content
"""
# Get the collection_id from the simulation
collection_id = get_asset_collection_id_for_simulation_id(sim_id=sim_id)
# Retrieve the asset collection
query_criteria = QueryCriteria().select_children('assets')
asset_collection = AssetCollection.get(id=collection_id, query_criteria=query_criteria)
# Return dictionary
ret = {}
# For each requested path, get the file content
for rpath in paths:
if remove_prefix and rpath.startswith(remove_prefix):
path = rpath[len(remove_prefix):]
else:
path = rpath
# Retrieve the relative_path and the file_name for the given path
relative_path, file_name = os.path.split(path)
relative_path = relative_path.strip('\\').strip('/')
# Look for the asset file in the collection
af = None
for asset_file in asset_collection.assets:
if asset_file.file_name == file_name and (asset_file.relative_path or '') == relative_path:
af = asset_file
break
# We did not find the asset in the collection -> error
if af is None:
raise Exception('Asset not found:\n%s %s \n%s' %
(relative_path, file_name, pretty_display_assets_from_collection(asset_collection.assets)))
# Retrieve the file
result = af.retrieve()
# write the file - result is written as output_directory/file_name, where file_name (with no pathing)
if output_directory:
output_file = os.path.normpath(os.path.join(output_directory, os.path.split(path)[1]))
dirname = os.path.dirname(output_file)
os.makedirs(dirname, exist_ok=True)
with open(output_file, 'wb') as f:
f.write(result)
# No matter what add to the return
ret[rpath] = result
return ret
def get_asset_collection_by_tag(tag_name, tag_value, query_criteria=None):
"""
Looks to see if a collection id exists for a given collection tag
:param collection_tag: An asset collection tag that uniquely identifies an asset collection
:return: An asset collection id if ONE match is found, else None (for 0 or 2+ matches)
"""
query_criteria = query_criteria or QueryCriteria().select_children('assets')
query_criteria.where_tag('%s=%s' % (tag_name, tag_value))
result = AssetCollection.get(query_criteria=query_criteria)
if len(result) >= 1: return result[0]
return None
def download_asset_collection(collection, output_folder):
if not isinstance(collection, AssetCollection):
collection = AssetCollection.get(collection, query_criteria=QueryCriteria().select_children('assets'))
# Get the files
if len(collection.assets) > 0:
# Download the collection as zip
zip_path = os.path.join(output_folder, 'temp.zip')
with open(zip_path, 'wb') as outfile:
outfile.write(collection.retrieve_as_zip())
# Extract it
zip_ref = zipfile.ZipFile(zip_path, 'r')
zip_ref.extractall(output_folder)
zip_ref.close()
# Delete the temporary zip
os.remove(zip_path)
def create_asset_collection(path_to_ac, name, other_tags=None):
"""
Create python asset collection from folder
"""
if other_tags is None:
other_tags = dict()
path_to_ac = os.path.normpath(path_to_ac)
if not os.path.exists(path_to_ac) or not os.path.isdir(path_to_ac):
print('Path \'{0}\' doesn\'t exist or is not a directory'.format(
path_to_ac))
exit()
tags = other_tags
tags['python_version'] = sys.version
tags['os'] = sys.platform
tags['Name'] = name
print("Adding files to asset Collection")
ac = AssetCollection()
ac.set_tags(tags)
for (dirpath, dirnames, filenames) in os.walk(path_to_ac):
for fn in filenames:
rp = os.path.relpath(dirpath,
path_to_ac) if dirpath != path_to_ac else ''
print('Adding {0}'.format(os.path.join(rp, fn)))
with open(os.path.join(dirpath, fn), 'rb') as fp:
acf = AssetCollectionFile(fn,
rp,
md5_checksum=hashlib.md5(
fp.read()).hexdigest())
ac.add_asset(acf, os.path.join(dirpath, fn))
print('Saving collection')
ac.save()
print('done - created AC ' + str(ac.id))
def generate_climate_files(self):
# see InputDataWorker for other work options
self.wo = InputDataWorker(
demographics_file_path=self.demographics_file_path,
wo_output_path=self.work_order_path,
project_info=self.climate_project,
start_year=self.start_year,
num_years=self.num_years,
resolution=self.resolution,
idRef=self.idRef)
# login to COMPS (if not already logged in) to submit climate files generation work order
self.wo.wo_2_json()
from COMPS.Data.WorkItem import WorkerOrPluginKey, WorkItemState
from COMPS.Data import QueryCriteria
from COMPS.Data import WorkItem, WorkItemFile
from COMPS.Data import AssetCollection
workerkey = WorkerOrPluginKey(name='InputDataWorker',
version='1.0.0.0_RELEASE')
wi = WorkItem('dtk-tools InputDataWorker WorkItem', workerkey,
SetupParser.get('environment'))
wi.set_tags({
'dtk-tools': None,
'WorkItem type': 'InputDataWorker dtk-tools'
})
with open(self.work_order_path, 'rb') as workorder_file:
# wi.AddWorkOrder(workorder_file.read())
wi.add_work_order(data=workorder_file.read())
with open(self.demographics_file_path, 'rb') as demog_file:
wi.add_file(WorkItemFile(
os.path.basename(self.demographics_file_path), 'Demographics',
''),
data=demog_file.read())
wi.save()
print("Created request for climate files generation.")
print("Commissioning...")
wi.commission()
while wi.state not in (WorkItemState.Succeeded, WorkItemState.Failed,
WorkItemState.Canceled):
om('Waiting for climate generation to complete (current state: ' +
str(wi.state) + ')',
style='flushed')
time.sleep(5)
wi.refresh()
print("Climate files SUCCESSFULLY generated")
# Get the collection with our files
collections = wi.get_related_asset_collections()
collection_id = collections[0].id
comps_collection = AssetCollection.get(
collection_id,
query_criteria=QueryCriteria().select_children('assets'))
# Get the files
if len(comps_collection.assets) > 0:
print("Found output files:")
for asset in comps_collection.assets:
print("- %s (%s)" % (asset.file_name, file_size(asset.length)))
print("\nDownloading to %s..." % self.climate_files_output_path)
# Download the collection
download_asset_collection(comps_collection,
self.climate_files_output_path)
# return filenames; this use of re in conjunction w/ glob is not great; consider refactor
rain_bin_re = os.path.abspath(self.climate_files_output_path +
'/*rain*.bin')
humidity_bin_re = os.path.abspath(self.climate_files_output_path +
'/*humidity*.bin')
temperature_bin_re = os.path.abspath(
self.climate_files_output_path + '/*temperature*.bin')
rain_file_name = os.path.basename(glob.glob(rain_bin_re)[0])
humidity_file_name = os.path.basename(
glob.glob(humidity_bin_re)[0])
temperature_file_name = os.path.basename(
glob.glob(temperature_bin_re)[0])
print('Climate files SUCCESSFULLY stored.')
return rain_file_name, temperature_file_name, humidity_file_name
else:
print('No output files found')