def file_extension_ok(file_path, white_list=[], black_list=['py', 'pyc']):
"""Determine whether a file has an acceptable extension
Args:
file_path (str): name of the file to examine
white_list ([str]): list of file types allowed (defaults to empty list)
black_list ([str]): list of file types rejected (defaults to ['py', 'pyc']). Ignored if white_list is not empty
Returns:
If file is a directory: True
If white_list non-empty: True if the file's extension matches any in the list, otherwise False
If white_list is empty: False if the file's extension matches any in black_list, otherwise True
"""
import os
if os.path.isdir(file_path):
return True
if len(white_list) > 0:
in_list = False
for ext in white_list:
in_list = in_list or pkio.has_file_extension(file_path, ext)
if not in_list:
return False
return True
for ext in black_list:
if pkio.has_file_extension(file_path, ext):
return False
return True
def api_importFile(simulation_type):
"""
Args:
simulation_type (str): which simulation type
Params:
file: file data
folder: where to import to
"""
import sirepo.importer
error = None
f = None
try:
f = flask.request.files.get('file')
if not f:
raise sirepo.util.Error('must supply a file')
req = http_request.parse_params(
filename=f.filename,
folder=flask.request.form.get('folder'),
id=flask.request.form.get('simulationId'),
template=True,
type=simulation_type,
)
req.file_stream = f.stream
req.import_file_arguments = flask.request.form.get('arguments', '')
def s(data):
data.models.simulation.folder = req.folder
data.models.simulation.isExample = False
return _save_new_and_reply(data)
if pkio.has_file_extension(req.filename, 'json'):
data = sirepo.importer.read_json(req.file_stream.read(), req.type)
#TODO(pjm): need a separate URI interface to importer, added exception for rs4pi for now
# (dicom input is normally a zip file)
elif pkio.has_file_extension(req.filename,
'zip') and req.type != 'rs4pi':
data = sirepo.importer.read_zip(req.file_stream, sim_type=req.type)
else:
if not hasattr(req.template, 'import_file'):
raise sirepo.util.Error('Only zip files are supported')
with simulation_db.tmp_dir() as d:
data = req.template.import_file(req, tmp_dir=d, reply_op=s)
if 'error' in data:
return http_reply.gen_json(data)
return s(data)
except werkzeug.exceptions.HTTPException:
raise
except sirepo.util.Reply:
raise
except Exception as e:
pkdlog('{}: exception: {}', f and f.filename, pkdexc())
#TODO(robnagler) security issue here. Really don't want to report errors to user
error = str(e.args) if hasattr(e, 'args') else str(e)
return http_reply.gen_json({
'error':
error if error else 'An unknown error occurred',
})
def create_archive(sim):
"""Zip up the json file and its dependencies
Args:
sim (PKDict): parsed request
Returns:
py.path.Local: zip file name
"""
if not pkio.has_file_extension(sim.filename, ('zip', 'html')):
raise sirepo.util.NotFound(
'unknown file type={}; expecting html or zip'.format(sim.filename)
)
with simulation_db.tmp_dir() as d:
want_zip = sim.filename.endswith('zip')
f, c = _create_zip(sim, want_python=want_zip, out_dir=d)
if want_zip:
t = 'application/zip'
else:
f, t = _create_html(f, c)
return sirepo.http_reply.gen_file_as_attachment(
f,
content_type=t,
filename=sim.filename,
)
def test_importer(import_req):
from pykern.pkcollections import PKDict
from sirepo.template import elegant
import sirepo.lib
for fn in pkio.sorted_glob(pkunit.data_dir().join('*')):
if not pkio.has_file_extension(fn, ('ele', 'lte')) \
or fn.basename.endswith('.ele.lte'):
continue
k = PKDict()
pkdlog('file={}', fn)
if fn.basename.startswith('deviance-'):
try:
data = elegant.import_file(import_req(fn))
except Exception as e:
k.actual = f'{e}\n'
else:
k.actual = 'did not raise exception'
elif fn.ext == '.lte':
data = elegant.import_file(import_req(fn))
data['models']['commands'] = []
g = elegant._Generate(data)
g.sim()
j = g.jinja_env
k.actual = j.rpn_variables + j.lattice
else:
f = sirepo.lib.Importer('elegant').parse_file(fn).write_files(pkunit.work_dir())
k.actual_path = f.commands
pkunit.file_eq(fn.basename + '.txt', **k)
def _dicom_files(dirname):
files = PKDict(
ctmap=PKDict(),
rtdose=None,
rtstruct=None,
position=None,
additional_rtdose=[],
)
for path in pkio.walk_tree(dirname):
if not pkio.has_file_extension(str(path), 'dcm'):
continue
v = pydicom.dcmread(str(path),
specific_tags=[
'SOPClassUID',
'InstanceNumber',
'PatientPosition',
])
if v.SOPClassUID == _DICOM_CLASS.CT_IMAGE:
files.ctmap[int(v.InstanceNumber)] = str(path)
files.position = v.PatientPosition
elif v.SOPClassUID == _DICOM_CLASS.RT_DOSE:
if files.rtdose:
files.additional_rtdose.append(str(path))
else:
files.rtdose = str(path)
elif v.SOPClassUID == _DICOM_CLASS.RT_STRUCT:
files.rtstruct = str(path)
assert files.rtdose and files.rtstruct, 'Missing RTSTRUCT and/or RTDOSE'
return files
def test_importer():
from pykern import pkcollections
from pykern import pkio
from pykern.pkunit import pkeq
from sirepo.template import elegant
with pkunit.save_chdir_work():
for fn in pkio.sorted_glob(pkunit.data_dir().join('*')):
if not pkio.has_file_extension(fn, ('ele', 'lte')) \
or fn.basename.endswith('ele.lte'):
continue
error = None
try:
data = elegant.import_file(FlaskRequest(fn))
except Exception as e:
pkdlog(pkdexc())
error = e.message
if error:
actual = error
else:
if pkio.has_file_extension(fn, 'lte'):
data['models']['commands'] = []
actual = '{}{}'.format(
elegant._generate_variables(data),
elegant.generate_lattice(
data,
elegant._build_filename_map(data),
elegant._build_beamline_map(data),
pkcollections.Dict(),
),
)
else:
data2 = elegant.import_file(FlaskRequest(
'{}.lte'.format(fn)),
test_data=data)
actual = elegant._generate_commands(
data2,
elegant._build_filename_map(data2),
elegant._build_beamline_map(data2),
pkcollections.Dict(),
)
outfile = fn.basename + '.txt'
pkio.write_text(outfile, actual)
expect = pkio.read_text(pkunit.data_dir().join(outfile))
#TODO(pjm): this takes too long if there are a lot of diffs
#assert expect == actual
pkeq(expect, actual)
def api_importFile(simulation_type=None):
"""
Args:
simulation_type (str): which simulation type
Params:
file: file data
folder: where to import to
"""
import sirepo.importer
error = None
f = None
try:
template = simulation_type and sirepo.template.import_module(
simulation_type)
f = flask.request.files.get('file')
assert f, \
ValueError('must supply a file')
if pkio.has_file_extension(f.filename, 'json'):
data = sirepo.importer.read_json(f.read(), template)
#TODO(pjm): need a separate URI interface to importer, added exception for rs4pi for now
# (dicom input is normally a zip file)
elif pkio.has_file_extension(f.filename,
'zip') and simulation_type != 'rs4pi':
data = sirepo.importer.read_zip(f.stream, template)
else:
assert simulation_type, \
'simulation_type is required param for non-zip|json imports'
assert hasattr(template, 'import_file'), \
ValueError('Only zip files are supported')
data = template.import_file(
flask.request,
simulation_db.simulation_lib_dir(simulation_type),
simulation_db.tmp_dir(),
)
#TODO(robnagler) need to validate folder
data.models.simulation.folder = flask.request.form['folder']
data.models.simulation.isExample = False
return _save_new_and_reply(data)
except Exception as e:
pkdlog('{}: exception: {}', f and f.filename, pkdexc())
error = str(e.message) if hasattr(e, 'message') else str(e)
return http_reply.gen_json({
'error':
error if error else 'An unknown error occurred',
})
def test_importer():
from pykern import pkcollections
from pykern import pkio
from pykern.pkunit import pkeq
from sirepo.template import elegant
with pkunit.save_chdir_work():
for fn in pkio.sorted_glob(pkunit.data_dir().join('*')):
if not pkio.has_file_extension(fn, ('ele', 'lte')) \
or fn.basename.endswith('ele.lte'):
continue
error = None
try:
data = elegant.import_file(FlaskRequest(fn))
except Exception as e:
pkdlog(pkdexc())
error = e.message
if error:
actual = error
else:
if pkio.has_file_extension(fn, 'lte'):
data['models']['commands'] = []
actual = '{}{}'.format(
elegant._generate_variables(data),
elegant.generate_lattice(
data,
elegant._build_filename_map(data),
elegant._build_beamline_map(data),
pkcollections.Dict(),
),
)
else:
data2 = elegant.import_file(FlaskRequest('{}.lte'.format(fn)), test_data=data)
actual = elegant._generate_commands(
data2,
elegant._build_filename_map(data2),
elegant._build_beamline_map(data2),
pkcollections.Dict(),
)
outfile = fn.basename + '.txt'
pkio.write_text(outfile, actual)
expect = pkio.read_text(pkunit.data_dir().join(outfile))
#TODO(pjm): this takes too long if there are a lot of diffs
#assert expect == actual
pkeq(expect, actual)
def test_importer(import_req):
from pykern.pkcollections import PKDict
from pykern.pkunit import pkeq
from sirepo.template import lattice
from sirepo.template import elegant
import sirepo.util
import flask
with pkunit.save_chdir_work():
for fn in pkio.sorted_glob(pkunit.data_dir().join('*')):
if not pkio.has_file_extension(fn, ('ele', 'lte')) \
or fn.basename.endswith('ele.lte'):
continue
error = None
try:
data = elegant.import_file(import_req(fn))
except Exception as e:
pkdlog(pkdexc())
error = str(e)
if error:
actual = error
else:
if pkio.has_file_extension(fn, 'lte'):
data['models']['commands'] = []
actual = '{}{}'.format(
elegant._generate_variables(data),
elegant._generate_lattice(
elegant._build_filename_map(data),
lattice.LatticeUtil(data, elegant._SCHEMA),
),
)
else:
#TODO(robnagler) test simulationId
data2 = elegant.import_file(import_req(fn.new(ext='ele.lte')), test_data=data)
actual = elegant._generate_commands(
elegant._build_filename_map(data2),
lattice.LatticeUtil(data2, elegant._SCHEMA),
)
outfile = fn.basename + '.txt'
pkio.write_text(outfile, actual)
expect = pkio.read_text(pkunit.data_dir().join(outfile))
pkeq(expect, actual)
def _extract_series_frames(simulation, dicom_dir):
#TODO(pjm): give user a choice between multiple study/series if present
selected_series = None
frames = {}
dicom_dose = None
rt_struct_path = None
res = {
'description': '',
}
for path in pkio.walk_tree(dicom_dir):
if pkio.has_file_extension(str(path), 'dcm'):
plan = dicom.read_file(str(path))
if plan.SOPClassUID == _DICOM_CLASS['RT_STRUCT']:
rt_struct_path = str(path)
elif plan.SOPClassUID == _DICOM_CLASS['RT_DOSE']:
res['dicom_dose'] = _summarize_rt_dose(simulation, plan)
plan.save_as(_dose_dicom_filename(simulation))
if plan.SOPClassUID != _DICOM_CLASS['CT_IMAGE']:
continue
orientation = _float_list(plan.ImageOrientationPatient)
if not (_EXPECTED_ORIENTATION == orientation).all():
continue
if not selected_series:
selected_series = plan.SeriesInstanceUID
res['StudyInstanceUID'] = plan.StudyInstanceUID
res['PixelSpacing'] = plan.PixelSpacing
if hasattr(plan, 'SeriesDescription'):
res['description'] = plan.SeriesDescription
if selected_series != plan.SeriesInstanceUID:
continue
info = {
'pixels': np.float32(plan.pixel_array),
'shape': plan.pixel_array.shape,
'ImagePositionPatient': _string_list(plan.ImagePositionPatient),
'ImageOrientationPatient': _float_list(plan.ImageOrientationPatient),
'PixelSpacing': _float_list(plan.PixelSpacing),
}
for f in ('FrameOfReferenceUID', 'StudyInstanceUID', 'SeriesInstanceUID', 'SOPInstanceUID'):
info[f] = getattr(plan, f)
z = _frame_id(info['ImagePositionPatient'][2])
info['frameId'] = z
if z in frames:
raise RuntimeError('duplicate frame with z coord: {}'.format(z))
_scale_pixel_data(plan, info['pixels'])
frames[z] = info
if not selected_series:
raise RuntimeError('No series found with {} orientation'.format(_EXPECTED_ORIENTATION))
if rt_struct_path:
res['regionsOfInterest'] = _summarize_rt_structure(simulation, dicom.read_file(rt_struct_path), frames.keys())
sorted_frames = []
res['frames'] = sorted_frames
for z in sorted(_float_list(frames.keys())):
sorted_frames.append(frames[_frame_id(z)])
return res
def _extract_series_frames(simulation, dicom_dir):
#TODO(pjm): give user a choice between multiple study/series if present
selected_series = None
frames = {}
dicom_dose = None
rt_struct_path = None
res = {
'description': '',
}
for path in pkio.walk_tree(dicom_dir):
if pkio.has_file_extension(str(path), 'dcm'):
plan = dicom.read_file(str(path))
if plan.SOPClassUID == _DICOM_CLASS['RT_STRUCT']:
rt_struct_path = str(path)
elif plan.SOPClassUID == _DICOM_CLASS['RT_DOSE']:
res['dicom_dose'] = _summarize_rt_dose(simulation, plan)
plan.save_as(_dose_dicom_filename(simulation))
if plan.SOPClassUID != _DICOM_CLASS['CT_IMAGE']:
continue
orientation = _float_list(plan.ImageOrientationPatient)
if not (_EXPECTED_ORIENTATION == orientation).all():
continue
if not selected_series:
selected_series = plan.SeriesInstanceUID
res['StudyInstanceUID'] = plan.StudyInstanceUID
res['PixelSpacing'] = plan.PixelSpacing
if hasattr(plan, 'SeriesDescription'):
res['description'] = plan.SeriesDescription
if selected_series != plan.SeriesInstanceUID:
continue
info = {
'pixels': np.float32(plan.pixel_array),
'shape': plan.pixel_array.shape,
'ImagePositionPatient': _string_list(plan.ImagePositionPatient),
'ImageOrientationPatient': _float_list(plan.ImageOrientationPatient),
'PixelSpacing': _float_list(plan.PixelSpacing),
}
for f in ('FrameOfReferenceUID', 'StudyInstanceUID', 'SeriesInstanceUID', 'SOPInstanceUID'):
info[f] = getattr(plan, f)
z = _frame_id(info['ImagePositionPatient'][2])
info['frameId'] = z
if z in frames:
raise RuntimeError('duplicate frame with z coord: {}'.format(z))
_scale_pixel_data(plan, info['pixels'])
frames[z] = info
if not selected_series:
raise RuntimeError('No series found with {} orientation'.format(_EXPECTED_ORIENTATION))
if rt_struct_path:
res['regionsOfInterest'] = _summarize_rt_structure(simulation, dicom.read_file(rt_struct_path), frames.keys())
sorted_frames = []
res['frames'] = sorted_frames
for z in sorted(_float_list(frames.keys())):
sorted_frames.append(frames[_frame_id(z)])
return res
def import_file(request, lib_dir=None, tmp_dir=None):
f = request.files['file']
filename = werkzeug.secure_filename(f.filename)
if not pkio.has_file_extension(str(filename), 'zip'):
raise RuntimeError('unsupported import filename: {}'.format(filename))
filepath = str(tmp_dir.join(_ZIP_FILE_NAME))
f.save(filepath)
data = simulation_db.default_data(SIM_TYPE)
data['models']['simulation']['name'] = filename
data['models']['simulation'][_TMP_INPUT_FILE_FIELD] = filepath
# more processing occurs below in prepare_for_client() after simulation dir is prepared
return data
def import_file(request, lib_dir=None, tmp_dir=None):
f = request.files['file']
filename = werkzeug.secure_filename(f.filename)
if not pkio.has_file_extension(str(filename), 'zip'):
raise RuntimeError('unsupported import filename: {}'.format(filename))
filepath = str(tmp_dir.join(_ZIP_FILE_NAME))
f.save(filepath)
data = simulation_db.default_data(SIM_TYPE)
data['models']['simulation']['name'] = filename
data['models']['simulation'][_TMP_INPUT_FILE_FIELD] = filepath
# more processing occurs below in prepare_for_client() after simulation dir is prepared
return data
def import_file(req, tmp_dir=None, **kwargs):
if not pkio.has_file_extension(req.filename, 'zip'):
raise sirepo.util.UserAlert('unsupported import filename: {}'.format(filename))
#TODO(pjm): writing to simulation lib for now, tmp_dir will get removed after this request
filepath = str(simulation_db.simulation_lib_dir(SIM_TYPE).join(_ZIP_FILE_NAME))
pkio.mkdir_parent_only(filepath)
with open(filepath, 'wb') as f:
f.write(req.file_stream.read())
data = simulation_db.default_data(SIM_TYPE)
data['models']['simulation']['name'] = req.filename
data['models']['simulation'][_TMP_INPUT_FILE_FIELD] = filepath
# more processing occurs in prepare_for_client() via:
# import_file => _save_new_and_reply => api_simulationData => prepare_for_client
return data
def api_importFile(simulation_type=None):
"""
Args:
simulation_type (str): which simulation type
Params:
file: file data
folder: where to import to
"""
import sirepo.importer
error = None
f = None
try:
template = simulation_type and sirepo.template.import_module(simulation_type)
f = flask.request.files.get('file')
assert f, \
ValueError('must supply a file')
if pkio.has_file_extension(f.filename, 'json'):
data = sirepo.importer.read_json(f.read(), template)
elif pkio.has_file_extension(f.filename, 'zip'):
data = sirepo.importer.read_zip(f.stream, template)
else:
assert simulation_type, \
'simulation_type is required param for non-zip|json imports'
data = template.import_file(
flask.request,
simulation_db.simulation_lib_dir(simulation_type),
simulation_db.tmp_dir(),
)
#TODO(robnagler) need to validate folder
data.models.simulation.folder = flask.request.form['folder']
return _save_new_and_reply(data.simulationType, data)
except Exception as e:
pkdlog('{}: exception: {}', f and f.filename, pkdexc())
error = e.message if hasattr(e, 'message') else str(e)
return _json_response({'error': error})
def create_archive(sim_type, sim_id, filename):
"""Zip up the json file and its dependencies
Args:
sim_type (str): simulation type
sim_id (str): simulation id
filename (str): for file type
Returns:
py.path.Local: zip file name
"""
from pykern import pkio
from sirepo import uri_router
if not pkio.has_file_extension(filename, ('zip', 'html')):
raise uri_router.NotFound(
'{}: unknown file type; expecting html or zip',
filename,
)
want_zip = filename.endswith('zip')
fp, data = _create_zip(sim_type, sim_id, want_python=want_zip)
if want_zip:
return fp, 'application/zip'
return _create_html(fp, data)
def create_archive(sim_type, sim_id, filename):
"""Zip up the json file and its dependencies
Args:
sim_type (str): simulation type
sim_id (str): simulation id
filename (str): for file type
Returns:
py.path.Local: zip file name
"""
from pykern import pkio
from sirepo import uri_router
if not pkio.has_file_extension(filename, ('zip', 'html')):
raise uri_router.NotFound(
'{}: unknown file type; expecting html or zip',
filename,
)
want_zip = filename.endswith('zip')
fp, data = _create_zip(sim_type, sim_id, want_python=want_zip)
if want_zip:
return fp, 'application/zip'
return _create_html(fp, data)
def test_has_file_extension():
from pykern.pkunit import pkeq
from pykern import pkio
pkeq(True, pkio.has_file_extension('x.ABC', 'abc'))
pkeq(True, pkio.has_file_extension(py.path.local('x.abc'), ('abc', 'def')))
