Commit 51bad4ac authored by Gabriel Couture's avatar Gabriel Couture

Fix typing errors

parent 4ae91cef
......@@ -30,36 +30,42 @@ def _format_dose_gradient(dose_gradient: str) -> str:
raise InvalidDoseGradientError("Dose Gradient must be {'AP', 'SI', 'all'}")
def get_rtstruct(geometry: str, spacing: float) -> List[pydicom.FileDataset]:
def get_rtstruct(geometry: str = 'all', spacing='all') -> List[pydicom.FileDataset]:
"""Get RTSTRUCT DICOM files.
Parameters
----------
geometry
Geometry: {'cones', 'spheres', 'cylinders'}.
Geometry: {'cones', 'spheres', 'cylinders', 'all'}.
spacing
Spacing in millimeters: {0.2, 1, 2, 3}.
Spacing in millimeters: {0.2, 1, 2, 3, 'all'}.
Returns
-------
List of DICOM file data set.
"""
directory, _ = os.path.split(__file__)
geometry_path = os.path.join(directory, 'data', 'RTSTRUCT', geometry)
if not os.path.exists(geometry_path):
raise WrongGeometryError("Geometry should be in {'cones', 'spheres', 'cylinders'}")
dicom_files = filter(
lambda i: f'{_format_spacing(spacing)}_' in i and 'X' not in i,
os.listdir(geometry_path)
)
dicom_paths = [os.path.join(geometry_path, i) for i in dicom_files]
if geometry == 'all':
geometry_path = os.path.join(directory, 'data', 'RTSTRUCT')
geometry_paths = [os.path.join(directory, 'data', 'RTSTRUCT', i) for i in os.listdir(geometry_path)]
else:
geometry_paths = [os.path.join(directory, 'data', 'RTSTRUCT', geometry)]
if False in [os.path.exists(i) for i in geometry_paths]:
raise WrongGeometryError("Geometry should be in {'cones', 'spheres', 'cylinders', 'all'}")
dicom_paths: List[str] = []
for geometry_path in geometry_paths:
dicom_files = filter(
lambda i: f'{_format_spacing(spacing)}_' in i or spacing == 'all',
os.listdir(geometry_path)
)
dicom_paths += [os.path.join(geometry_path, i) for i in dicom_files]
return [pydicom.dcmread(i) for i in dicom_paths]
def get_rtdose(dose_gradient: str, spacing: float) -> List[pydicom.FileDataset]:
def get_rtdose(dose_gradient: str = 'all', spacing='all') -> List[pydicom.FileDataset]:
"""Get RTDOSE DICOM files.
Parameters
......@@ -68,7 +74,7 @@ def get_rtdose(dose_gradient: str, spacing: float) -> List[pydicom.FileDataset]:
Dose gradient: {'AP', 'SI', 'all'}.
AP := Anterior Posterior, SI := Superior Inferior.
spacing
Spacing in millimeters: {1, 2, 3}.
Spacing in millimeters: {1, 2, 3, 'all'}.
Returns
-------
......@@ -78,8 +84,11 @@ def get_rtdose(dose_gradient: str, spacing: float) -> List[pydicom.FileDataset]:
rtdose_path = os.path.join(directory, 'data', 'RTDOSE')
dicom_files = filter(
lambda i: f'{spacing}mm'.replace('.0', '') in i and _format_dose_gradient(dose_gradient) in i,
os.listdir(rtdose_path)
lambda i: _format_dose_gradient(dose_gradient) in i or dose_gradient == 'all',
filter(
lambda i: f'{spacing}mm'.replace('.0', '') in i or spacing == 'all',
os.listdir(rtdose_path)
)
)
dicom_paths = [os.path.join(rtdose_path, i) for i in dicom_files]
......@@ -97,4 +106,3 @@ def get_rtplan() -> pydicom.FileDataset:
rtplan_path = os.path.join(directory, 'data', 'RTPLAN', 'RTPLAN.dcm')
return pydicom.dcmread(rtplan_path)
......@@ -20,8 +20,15 @@ class TestDao(unittest.TestCase):
result = dao.get_rtstruct(A_GEOMETRY, A_SPACING)
self.assertEqual(len(result), 1)
self.assertIsInstance(result[0], pydicom.FileDataset)
self.assertEqual(result[0]['Modality'].value, 'RTSTRUCT')
for i in result:
self.assert_is_rtstruct(i)
def test_givenAGeometry_whenGettingRtstruct_thenResultIsRtstructFileDataSetsWithAnySpacing(self):
result = dao.get_rtstruct(A_GEOMETRY)
self.assertEqual(len(result), 4)
for i in result:
self.assert_is_rtstruct(i)
def test_givenAGeometryAndASpacing_whenGettingRtstruct_thenRaiseWrongGeometryError(self):
self.assertRaises(
......@@ -33,8 +40,8 @@ class TestDao(unittest.TestCase):
result = dao.get_rtdose(A_DOSE_GRADIENTS, A_SPACING)
self.assertEqual(len(result), 1)
self.assertIsInstance(result[0], pydicom.FileDataset)
self.assertEqual(result[0]['Modality'].value, 'RTDOSE')
for i in result:
self.assert_is_rtdose(i)
def test_givenADoseGradientAndASpacing_whenGettingRtstruct_thenRaiseWrongGeometryError(self):
self.assertRaises(
......@@ -48,3 +55,10 @@ class TestDao(unittest.TestCase):
self.assertIsInstance(result, pydicom.FileDataset)
self.assertEqual(result['Modality'].value, 'RTPLAN')
def assert_is_rtstruct(self, data_set: pydicom.FileDataset):
self.assertIsInstance(data_set, pydicom.FileDataset)
self.assertEqual(data_set['Modality'].value, 'RTSTRUCT')
def assert_is_rtdose(self, data_set: pydicom.FileDataset):
self.assertIsInstance(data_set, pydicom.FileDataset)
self.assertEqual(data_set['Modality'].value, 'RTDOSE')
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