<commit_before>import os, string
from os import path
from glob import glob
from nipype.interfaces.base import (TraitedSpec, 
                                    DynamicTraitedSpec, 
                                    InputMultiPath, 
                                    File,
                                    Directory, 
                                    traits, 
                                    BaseInterface,
                                   )
import dicom
import dcmstack
from dcmstack.dcmmeta import NiftiWrapper
import nibabel as nb

def sanitize_path_comp(path_comp):
    result = []
    for char in path_comp:
        if not char in string.letters + string.digits + '-_.':
            result.append('_')
        else:
            result.append(char)
    return ''.join(result)

class NiftiGeneratorBaseInputSpec(TraitedSpec):
    out_format = traits.Str(desc="String which can be formatted with "
                            "meta data to create the output filename(s)")
    out_ext = traits.Str('.nii.gz', 
                         usedefault=True, 
                         desc="Determines output file type")

class NiftiGeneratorBase(BaseInterface):
    def _get_out_path(self, meta):
        if self.inputs.out_format:
            out_fmt = self.inputs.out_format
        else:
            out_fmt = []
            if 'SeriesNumber' in meta:
                out_fmt.append('%(SeriesNumber)03d')
            if 'ProtocolName' in meta:
                out_fmt.append('%(ProtocolName)s')
            elif 'SeriesDescription' in meta:
                out_fmt.append('%(SeriesDescription)s')
            else:
                out_fmt.append('sequence')
            out_fmt = '-'.join(out_fmt)
        out_fn = (out_fmt % meta) + self.inputs.out_ext
        out_fn = sanitize_path_comp(out_fn)
        return path.join(os.getcwd(), out_fn)

class DcmStackInputSpec(NiftiGeneratorBaseInputSpec):
    dicom_files = traits.Either(InputMultiPath(File(exists=True)),
                                Directory(exists=True),
                                traits.Str(), 
                                mandatory=True)
    embed_meta = traits.Bool(desc="Embed DICOM meta data into result")
    
class DcmStackOutputSpec(TraitedSpec):
    out_file = traits.File(exists=True)

class DcmStack(NiftiGeneratorBase):
    '''Create one Nifti file from a set of DICOM files'''
    input_spec = DcmStackInputSpec
    output_spec = DcmStackOutputSpec
    
    def _get_filelist(self, trait_input):
        if isinstance(trait_input, str):
            if path.isdir(trait_input):
                return glob(path.join(trait_input, '*.dcm'))
            else:
                return glob(trait_input)
        
        return trait_input
    
    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        stack = dcmstack.DicomStack()
        for src_path in src_paths:
            src_dcm = dicom.read_file(src_path)
            stack.add_dcm(src_dcm)
        nii = stack.to_nifti(embed_meta=True)
        nw = NiftiWrapper(nii)
        self.out_path = self._get_out_path(nw.meta_ext.get_class_dict(('global', 'const')))
        if not self.inputs.embed_meta:
            nw.remove_extension()
        nb.save(nii, self.out_path)
        return runtime
        
    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_file"] = self.out_path
        return outputs

class GroupAndStackOutputSpec(TraitedSpec):
    out_list = traits.List(desc="List of output nifti files")

class GroupAndStack(NiftiGeneratorBase):
    '''Create (potentially) multiple Nifti files for a set of DICOM 
    files.'''
    input_spec = DcmStackInputSpec
    output_spec = GroupAndStackOutputSpec
    
    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        stacks = \<commit_msg>ENH: Default to 'force=True' when reading DICOM files.<commit_after>import os, string
from os import path
from glob import glob
from nipype.interfaces.base import (TraitedSpec, 
                                    DynamicTraitedSpec, 
                                    InputMultiPath, 
                                    File,
                                    Directory, 
                                    traits, 
                                    BaseInterface,
                                   )
import dicom
import dcmstack
from dcmstack.dcmmeta import NiftiWrapper
import nibabel as nb

def sanitize_path_comp(path_comp):
    result = []
    for char in path_comp:
        if not char in string.letters + string.digits + '-_.':
            result.append('_')
        else:
            result.append(char)
    return ''.join(result)

class NiftiGeneratorBaseInputSpec(TraitedSpec):
    out_format = traits.Str(desc="String which can be formatted with "
                            "meta data to create the output filename(s)")
    out_ext = traits.Str('.nii.gz', 
                         usedefault=True, 
                         desc="Determines output file type")

class NiftiGeneratorBase(BaseInterface):
    def _get_out_path(self, meta):
        if self.inputs.out_format:
            out_fmt = self.inputs.out_format
        else:
            out_fmt = []
            if 'SeriesNumber' in meta:
                out_fmt.append('%(SeriesNumber)03d')
            if 'ProtocolName' in meta:
                out_fmt.append('%(ProtocolName)s')
            elif 'SeriesDescription' in meta:
                out_fmt.append('%(SeriesDescription)s')
            else:
                out_fmt.append('sequence')
            out_fmt = '-'.join(out_fmt)
        out_fn = (out_fmt % meta) + self.inputs.out_ext
        out_fn = sanitize_path_comp(out_fn)
        return path.join(os.getcwd(), out_fn)

class DcmStackInputSpec(NiftiGeneratorBaseInputSpec):
    dicom_files = traits.Either(InputMultiPath(File(exists=True)),
                                Directory(exists=True),
                                traits.Str(), 
                                mandatory=True)
    embed_meta = traits.Bool(desc="Embed DICOM meta data into result")
    
class DcmStackOutputSpec(TraitedSpec):
    out_file = traits.File(exists=True)

class DcmStack(NiftiGeneratorBase):
    '''Create one Nifti file from a set of DICOM files'''
    input_spec = DcmStackInputSpec
    output_spec = DcmStackOutputSpec
    
    def _get_filelist(self, trait_input):
        if isinstance(trait_input, str):
            if path.isdir(trait_input):
                return glob(path.join(trait_input, '*.dcm'))
            else:
                return glob(trait_input)
        
        return trait_input
    
    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        stack = dcmstack.DicomStack()
        for src_path in src_paths:
            src_dcm = dicom.read_file(src_path, force=True)
            stack.add_dcm(src_dcm)
        nii = stack.to_nifti(embed_meta=True)
        nw = NiftiWrapper(nii)
        self.out_path = self._get_out_path(nw.meta_ext.get_class_dict(('global', 'const')))
        if not self.inputs.embed_meta:
            nw.remove_extension()
        nb.save(nii, self.out_path)
        return runtime
        
    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_file"] = self.out_path
        return outputs

class GroupAndStackOutputSpec(TraitedSpec):
    out_list = traits.List(desc="List of output nifti files")

class GroupAndStack(NiftiGeneratorBase):
    '''Create (potentially) multiple Nifti files for a set of DICOM 
    files.'''
    input_spec = DcmStackInputSpec
    output_spec = GroupAndStackOutputSpec
    
    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        stacks = \
{'fraction_non_alphanumeric': 0.08659609743655146, 'fraction_numerical': 0.0005101390128810101, 'mean_word_length': 1.921758569299553, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5502272', 'n_tokens_mistral': 2306, 'n_tokens_neox': 2228, 'n_words': 470}