<commit_before>from numba.core.typing.templates import ConcreteTemplate
from numba.core import types, typing, funcdesc, config, compiler
from numba.core.compiler import (sanitize_compile_result_entries, CompilerBase,
                                 DefaultPassBuilder, Flags, Option,
                                 CompileResult)
from numba.core.compiler_lock import global_compiler_lock
from numba.core.compiler_machinery import (LoweringPass, AnalysisPass,
                                           PassManager, register_pass)
from numba.core.errors import NumbaInvalidConfigWarning, TypingError
from numba.core.typed_passes import (IRLegalization, NativeLowering,
                                     AnnotateTypes)
from warnings import warn
from numba.cuda.api import get_current_device

def _nvvm_options_type(x):
    if x is None:
        return None

    else:
        assert isinstance(x, dict)
        return x

class CUDAFlags(Flags):
    nvvm_options = Option(
        type=_nvvm_options_type,
        default=None,
        doc="NVVM options",
    )

# The CUDACompileResult (CCR) has a specially-defined entry point equal to its
# id.  This is because the entry point is used as a key into a dict of
# overloads by the base dispatcher. The id of the CCR is the only small and
# unique property of a CompileResult in the CUDA target (cf. the CPU target,
# which uses its entry_point, which is a pointer value).
#
# This does feel a little hackish, and there are two ways in which this could
# be improved:
#
# 1. We could change the core of Numba so that each CompileResult has its own
#    unique ID that can be used as a key - e.g. a count, similar to the way in
#    which types have unique counts.
# 2. At some future time when kernel launch uses a compiled function, the entry
#    point will no longer need to be a synthetic value, but will instead be a
#    pointer to the compiled function as in the CPU target.

class CUDACompileResult(CompileResult):
    @property
    def entry_point(self):
        return id(self)

def cuda_compile_result(**entries):
    entries = sanitize_compile_result_entries(entries)
    return CUDACompileResult(**entries)

@register_pass(mutates_CFG=True, analysis_only=False)
class CUDABackend(LoweringPass):

    _name = "cuda_backend"

    def __init__(self):
        LoweringPass.__init__(self)

    def run_pass(self, state):
        """
        Back-end: Packages lowering output in a compile result
        """
        lowered = state['cr']
        signature = typing.signature(state.return_type, *state.args)

        state.cr = cuda_compile_result(
            typing_context=state.typingctx,
            target_context=state.targetctx,
            typing_error=state.status.fail_reason,
            type_annotation=state.type_annotation,
            library=state.library,
            call_helper=lowered.call_helper,
            signature=signature,
            fndesc=lowered.fndesc,
        )
        return True

@register_pass(mutates_CFG=False, analysis_only=False)
class CreateLibrary(LoweringPass):
    """
    Create a CUDACodeLibrary for the NativeLowering pass to populate. The
    NativeLowering pass will create a code library if none exists, but we need
    to set it up with nvvm_options from the flags if they are present.
    """

    _name = "create_library"

    def __init__(self):
        LoweringPass.__init__(self)

    def run_pass(self, state):
        codegen = state.targetctx.codegen()
        name = state.func_id.func_qualname
        nvvm_options = state.flags.nvvm_options
        state.library = codegen.create_library(name, nvvm_options=nvvm_options)
        # Enable object caching upfront so that the library can be serialized.
        state.library.enable_object_caching()

        return True

@register_pass(mutates_CFG=False, analysis_only=True)
class CUDALegalization(AnalysisPass):

    _name = "cuda_legalization"

    def __init__(self):
        AnalysisPass.__init__(self)

    def run_pass(self, state):
        # Early return if NVVM 7
        from numba.cuda.cudadrv.nvvm import NVVM
        if NVVM().is_nvvm70:
            return False
        # NVVM < 7, need to check for charseq
        typmap = state.typemap

        def check_dtype(dtype):
            if isinstance(dtype, (types.UnicodeCharSeq, types.CharSeq)):
                msg = (f"{k} is a char sequence type. This type is not "
                       "supported with CUDA toolkit versions < 11.2. To "
                       "use this type, you need to update your CUDA "
                       "toolkit - try 'conda install cudatoolkit=11' if "
                       "you are using conda to manage your environment.")
                raise TypingError(msg)
            elif isinstance(dtype, types.Record):
                for subdtype in dtype.fields.items():
                    # subdtype is a (name, _RecordField) pair
                    check_dtype(subdtype[1].type)

        for k, v in typmap.items():
            if isinstance(v, types.Array):
                check_dtype(v.dtype)
        return False

class CUDACompiler(CompilerBase):
    def define_pipelines(self):
        dpb = DefaultPassBuilder
        pm = PassManager('cuda')

        untyped_passes = dpb.define_untyped_pipeline(self.state)
        pm.passes.extend(untyped_passes.passes)

        typed_passes = dpb.define_typed_pipeline(self.state)
        pm.passes.extend(typed_passes.passes)
        pm.add_pass(CUDALegalization, "CUDA legalization")

        lowering_passes = self.define_cuda_lowering_pipeline(self.state)
        pm.passes.extend(lowering_passes.passes)

        pm.finalize()
        return [pm]

    def define_cuda_lowering_pipeline(self, state):
        pm = PassManager('cuda_lowering')
        # legalise
        pm.add_pass(IRLegalization,
                    "ensure IR is legal prior to lowering")
        pm.add_pass(AnnotateTypes, "annotate types")

        # lower
        pm.add_pass(CreateLibrary, "create library")
        pm.add_pass(NativeLowering, "native lowering")
        pm.add_pass(CUDABackend, "cuda backend")

        pm.finalize()
        return pm

@global_compiler_lock
def compile_cuda(pyfunc, return_type, args, debug=False, lineinfo=False,
                 inline=False, fastmath=False, nvvm_options=None):
    from .descriptor import cuda_target
    typingctx = cuda_target.typing_context
    targetctx = cuda_target.target_context

    flags = CUDAFlags()
    # Do not compile (generate native code), just lower (to LLVM)
    flags.no_compile = True
    flags.no_cpython_wrapper = True
    flags.no_cfunc_wrapper = True
    if debug or lineinfo:
        # Note both debug and lineinfo turn on debug information in the
        # compiled code, but we keep them separate arguments in case we
        # later want to overload some other behavior on the debug flag.
        # In particular, -opt=3 is not supported with -g.
        flags.debuginfo = True
        flags.error_model = 'python'
    else:
        flags.error_model = 'numpy'
    if inline:
        flags.forceinline = True
    if fastmath:
        flags.fastmath = True
    if nvvm_options:
        flags.nvvm_options = nvvm_options

    # Run compilation pipeline
    cres = compiler.compile_extra(typingctx=typingctx,
                                  targetctx=targetctx,
                                  func=pyfunc,
                                  args=args,
                                  return_type=return_type,
                                  flags=flags,
                                  locals={},
                                  pipeline_class=CUDACompiler)

    library = cres.library
    library.finalize()

    return cres

@global_compiler_lock
def compile_ptx(pyfunc, args, debug=False, lineinfo=False, device=False,
                fastmath=False, cc=None, opt=True):
    """Compile a Python function to PTX for a given set of argument types.

    :param pyfunc: The Python function to compile.
    :param args: A tuple of argument types to compile for.
    :param debug: Whether to include debug info in the generated PTX.
    :type debug: bool
    :param lineinfo: Whether to include a line mapping from the generated PTX
                     to the source code. Usually this is used with optimized
                     code (since debug mode would automatically include this),
                     so we want debug info in the LLVM but only the line
                     mapping in the final PTX.
    :type lineinfo: bool
    :param device: Whether to compile a device function. Defaults to ``False``,
                   to compile global kernel functions.
    :type device: bool
    :param fastmath: Whether to enable fast math flags (ftz=1, prec_sqrt=0,
                     prec_div=, and fma=1)
    :type fastmath: bool
    :param cc: Compute capability to compile for, as a tuple ``(MAJOR, MINOR)``.
               Defaults to ``(5, 3)``.
    :type cc: tuple
    :param opt: Enable optimizations. Defaults to ``True``.
    :type opt: bool
    :return: (ptx, resty): The PTX code and inferred return type
    :rtype: tuple
    """
    if debug and opt:
        msg = ("debug=True with opt=True (the default) "
               "is not supported by CUDA. This may result in a crash"
               " - set debug=False or opt=False.")
        warn(NumbaInvalidConfigWarning(msg))

    nvvm_options = {
        'debug': debug,
        'lineinfo': lineinfo,
        'fastmath': fastmath,
        'opt': 3 if opt else 0
    }

    cres = compile_cuda(pyfunc, None, args, debug=debug, lineinfo=lineinfo,
                        fastmath=fastmath,
                        nvvm_options=nvvm_options)
    resty = cres.signature.return_type
    if device:
        lib = cres.library
    else:
        tgt = cres.target_context
        filename = cres.type_annotation.filename
        linenum = int(cres.type_annotation.linenum)
        lib, kernel = tgt.prepare_cuda_kernel(cres.library, cres.fndesc, debug,
                                              nvvm_options, filename, linenum)

    cc = cc or config.CUDA_DEFAULT_PTX_CC
    ptx = lib.get_asm_str(cc=cc)
    return ptx, resty

def compile_ptx_for_current_device(pyfunc, args, debug=False, lineinfo=False,
                                   device=False, fastmath=False, opt=True):
    """Compile a Python function to PTX for a given set of argument types for
    the current device's compute capabilility. This calls :func:`compile_ptx`
    with an appropriate ``cc`` value for the current device."""
    cc = get_current_device().compute_capability
    return compile_ptx(pyfunc, args, debug=debug, lineinfo=lineinfo,
                       device=device, fastmath=fastmath, cc=cc, opt=True)

def declare_device_function(name, restype, argtypes):
    from .descriptor import cuda_target
    typingctx = cuda_target.typing_context
    targetctx = cuda_target.target_context
    sig = typing.signature(restype, *argtypes)
    extfn = ExternFunction(name, sig)

    class device_function_template(ConcreteTemplate):
        key = extfn
        cases = [sig]

    fndesc = funcdesc.ExternalFunctionDescriptor(
        name=name, restype=restype, argtypes=argtypes)
    typingctx.insert_user_function(extfn, device_function_template)
    targetctx.insert_user_function(extfn, fndesc)
    return extfn

class ExternFunction(object):
    def __init__(self, name, sig):
        self.name = name
        self.sig = sig
<commit_msg>Get lineno and filename from pyfunc instead of type annotation<commit_after>from numba.core.typing.templates import ConcreteTemplate
from numba.core import types, typing, funcdesc, config, compiler
from numba.core.compiler import (sanitize_compile_result_entries, CompilerBase,
                                 DefaultPassBuilder, Flags, Option,
                                 CompileResult)
from numba.core.compiler_lock import global_compiler_lock
from numba.core.compiler_machinery import (LoweringPass, AnalysisPass,
                                           PassManager, register_pass)
from numba.core.errors import NumbaInvalidConfigWarning, TypingError
from numba.core.typed_passes import (IRLegalization, NativeLowering,
                                     AnnotateTypes)
from warnings import warn
from numba.cuda.api import get_current_device
from numba.core.bytecode import FunctionIdentity

def _nvvm_options_type(x):
    if x is None:
        return None

    else:
        assert isinstance(x, dict)
        return x

class CUDAFlags(Flags):
    nvvm_options = Option(
        type=_nvvm_options_type,
        default=None,
        doc="NVVM options",
    )

# The CUDACompileResult (CCR) has a specially-defined entry point equal to its
# id.  This is because the entry point is used as a key into a dict of
# overloads by the base dispatcher. The id of the CCR is the only small and
# unique property of a CompileResult in the CUDA target (cf. the CPU target,
# which uses its entry_point, which is a pointer value).
#
# This does feel a little hackish, and there are two ways in which this could
# be improved:
#
# 1. We could change the core of Numba so that each CompileResult has its own
#    unique ID that can be used as a key - e.g. a count, similar to the way in
#    which types have unique counts.
# 2. At some future time when kernel launch uses a compiled function, the entry
#    point will no longer need to be a synthetic value, but will instead be a
#    pointer to the compiled function as in the CPU target.

class CUDACompileResult(CompileResult):
    @property
    def entry_point(self):
        return id(self)

def cuda_compile_result(**entries):
    entries = sanitize_compile_result_entries(entries)
    return CUDACompileResult(**entries)

@register_pass(mutates_CFG=True, analysis_only=False)
class CUDABackend(LoweringPass):

    _name = "cuda_backend"

    def __init__(self):
        LoweringPass.__init__(self)

    def run_pass(self, state):
        """
        Back-end: Packages lowering output in a compile result
        """
        lowered = state['cr']
        signature = typing.signature(state.return_type, *state.args)

        state.cr = cuda_compile_result(
            typing_context=state.typingctx,
            target_context=state.targetctx,
            typing_error=state.status.fail_reason,
            type_annotation=state.type_annotation,
            library=state.library,
            call_helper=lowered.call_helper,
            signature=signature,
            fndesc=lowered.fndesc,
        )
        return True

@register_pass(mutates_CFG=False, analysis_only=False)
class CreateLibrary(LoweringPass):
    """
    Create a CUDACodeLibrary for the NativeLowering pass to populate. The
    NativeLowering pass will create a code library if none exists, but we need
    to set it up with nvvm_options from the flags if they are present.
    """

    _name = "create_library"

    def __init__(self):
        LoweringPass.__init__(self)

    def run_pass(self, state):
        codegen = state.targetctx.codegen()
        name = state.func_id.func_qualname
        nvvm_options = state.flags.nvvm_options
        state.library = codegen.create_library(name, nvvm_options=nvvm_options)
        # Enable object caching upfront so that the library can be serialized.
        state.library.enable_object_caching()

        return True

@register_pass(mutates_CFG=False, analysis_only=True)
class CUDALegalization(AnalysisPass):

    _name = "cuda_legalization"

    def __init__(self):
        AnalysisPass.__init__(self)

    def run_pass(self, state):
        # Early return if NVVM 7
        from numba.cuda.cudadrv.nvvm import NVVM
        if NVVM().is_nvvm70:
            return False
        # NVVM < 7, need to check for charseq
        typmap = state.typemap

        def check_dtype(dtype):
            if isinstance(dtype, (types.UnicodeCharSeq, types.CharSeq)):
                msg = (f"{k} is a char sequence type. This type is not "
                       "supported with CUDA toolkit versions < 11.2. To "
                       "use this type, you need to update your CUDA "
                       "toolkit - try 'conda install cudatoolkit=11' if "
                       "you are using conda to manage your environment.")
                raise TypingError(msg)
            elif isinstance(dtype, types.Record):
                for subdtype in dtype.fields.items():
                    # subdtype is a (name, _RecordField) pair
                    check_dtype(subdtype[1].type)

        for k, v in typmap.items():
            if isinstance(v, types.Array):
                check_dtype(v.dtype)
        return False

class CUDACompiler(CompilerBase):
    def define_pipelines(self):
        dpb = DefaultPassBuilder
        pm = PassManager('cuda')

        untyped_passes = dpb.define_untyped_pipeline(self.state)
        pm.passes.extend(untyped_passes.passes)

        typed_passes = dpb.define_typed_pipeline(self.state)
        pm.passes.extend(typed_passes.passes)
        pm.add_pass(CUDALegalization, "CUDA legalization")

        lowering_passes = self.define_cuda_lowering_pipeline(self.state)
        pm.passes.extend(lowering_passes.passes)

        pm.finalize()
        return [pm]

    def define_cuda_lowering_pipeline(self, state):
        pm = PassManager('cuda_lowering')
        # legalise
        pm.add_pass(IRLegalization,
                    "ensure IR is legal prior to lowering")
        pm.add_pass(AnnotateTypes, "annotate types")

        # lower
        pm.add_pass(CreateLibrary, "create library")
        pm.add_pass(NativeLowering, "native lowering")
        pm.add_pass(CUDABackend, "cuda backend")

        pm.finalize()
        return pm

@global_compiler_lock
def compile_cuda(pyfunc, return_type, args, debug=False, lineinfo=False,
                 inline=False, fastmath=False, nvvm_options=None):
    from .descriptor import cuda_target
    typingctx = cuda_target.typing_context
    targetctx = cuda_target.target_context

    flags = CUDAFlags()
    # Do not compile (generate native code), just lower (to LLVM)
    flags.no_compile = True
    flags.no_cpython_wrapper = True
    flags.no_cfunc_wrapper = True
    if debug or lineinfo:
        # Note both debug and lineinfo turn on debug information in the
        # compiled code, but we keep them separate arguments in case we
        # later want to overload some other behavior on the debug flag.
        # In particular, -opt=3 is not supported with -g.
        flags.debuginfo = True
        flags.error_model = 'python'
    else:
        flags.error_model = 'numpy'
    if inline:
        flags.forceinline = True
    if fastmath:
        flags.fastmath = True
    if nvvm_options:
        flags.nvvm_options = nvvm_options

    # Run compilation pipeline
    cres = compiler.compile_extra(typingctx=typingctx,
                                  targetctx=targetctx,
                                  func=pyfunc,
                                  args=args,
                                  return_type=return_type,
                                  flags=flags,
                                  locals={},
                                  pipeline_class=CUDACompiler)

    library = cres.library
    library.finalize()

    return cres

@global_compiler_lock
def compile_ptx(pyfunc, args, debug=False, lineinfo=False, device=False,
                fastmath=False, cc=None, opt=True):
    """Compile a Python function to PTX for a given set of argument types.

    :param pyfunc: The Python function to compile.
    :param args: A tuple of argument types to compile for.
    :param debug: Whether to include debug info in the generated PTX.
    :type debug: bool
    :param lineinfo: Whether to include a line mapping from the generated PTX
                     to the source code. Usually this is used with optimized
                     code (since debug mode would automatically include this),
                     so we want debug info in the LLVM but only the line
                     mapping in the final PTX.
    :type lineinfo: bool
    :param device: Whether to compile a device function. Defaults to ``False``,
                   to compile global kernel functions.
    :type device: bool
    :param fastmath: Whether to enable fast math flags (ftz=1, prec_sqrt=0,
                     prec_div=, and fma=1)
    :type fastmath: bool
    :param cc: Compute capability to compile for, as a tuple ``(MAJOR, MINOR)``.
               Defaults to ``(5, 3)``.
    :type cc: tuple
    :param opt: Enable optimizations. Defaults to ``True``.
    :type opt: bool
    :return: (ptx, resty): The PTX code and inferred return type
    :rtype: tuple
    """
    if debug and opt:
        msg = ("debug=True with opt=True (the default) "
               "is not supported by CUDA. This may result in a crash"
               " - set debug=False or opt=False.")
        warn(NumbaInvalidConfigWarning(msg))

    nvvm_options = {
        'debug': debug,
        'lineinfo': lineinfo,
        'fastmath': fastmath,
        'opt': 3 if opt else 0
    }

    cres = compile_cuda(pyfunc, None, args, debug=debug, lineinfo=lineinfo,
                        fastmath=fastmath,
                        nvvm_options=nvvm_options)
    resty = cres.signature.return_type
    if device:
        lib = cres.library
    else:
        tgt = cres.target_context
        func_identity = FunctionIdentity.from_function(pyfunc)
        filename = func_identity.filename
        linenum = func_identity.firstlineno
        lib, kernel = tgt.prepare_cuda_kernel(cres.library, cres.fndesc, debug,
                                              nvvm_options, filename, linenum)

    cc = cc or config.CUDA_DEFAULT_PTX_CC
    ptx = lib.get_asm_str(cc=cc)
    return ptx, resty

def compile_ptx_for_current_device(pyfunc, args, debug=False, lineinfo=False,
                                   device=False, fastmath=False, opt=True):
    """Compile a Python function to PTX for a given set of argument types for
    the current device's compute capabilility. This calls :func:`compile_ptx`
    with an appropriate ``cc`` value for the current device."""
    cc = get_current_device().compute_capability
    return compile_ptx(pyfunc, args, debug=debug, lineinfo=lineinfo,
                       device=device, fastmath=fastmath, cc=cc, opt=True)

def declare_device_function(name, restype, argtypes):
    from .descriptor import cuda_target
    typingctx = cuda_target.typing_context
    targetctx = cuda_target.target_context
    sig = typing.signature(restype, *argtypes)
    extfn = ExternFunction(name, sig)

    class device_function_template(ConcreteTemplate):
        key = extfn
        cases = [sig]

    fndesc = funcdesc.ExternalFunctionDescriptor(
        name=name, restype=restype, argtypes=argtypes)
    typingctx.insert_user_function(extfn, device_function_template)
    targetctx.insert_user_function(extfn, fndesc)
    return extfn

class ExternFunction(object):
    def __init__(self, name, sig):
        self.name = name
        self.sig = sig
<|endoftext|>


<commit_before># MAEC 4.0.1 to MAEC 4.1 Converter Script
# Translates a MAEC 4.0.1 Package or Bundle into a valid MAEC 4.1 Package or Bundle

import sys
import maec
from maec.bundle.bundle import Bundle
from maec.package.package import Package

# Update the MAEC v4.0.1 file to MAEC v4.1
def update_maec(infilename, outfilename):
    # Parse the input document using the parse_xml_instance() method
    maec_objects = maec.parse_xml_instance(infilename, check_version = False)

    # Get the API Object from the parsed input
    api_object = maec_objects['api']

    # Determine if we're dealing with a Package or Bundle
    if isinstance(api_object, Package):
        # Update the Package schema_version
        api_object.schema_version = "2.1"
        for malware_subject in api_object.malware_subjects:
            for analysis in malware_subject.analyses:
                # Replace the Analysis type value of "manual" with "in-depth"
                if analysis.type and analysis.type == "manual":
                    analysis.type = "in-depth"
            # Update the schema_versions on the Bundles
            for bundle in malware_subject.findings_bundles.bundles:
                bundle.schema_version = "4.1"
    elif isinstance(api_object, Bundle):
        # Update the Bundle schema_version
        api_object.schema_version = "4.1"

    # Output the updated MAEC object to XML
    api_object.to_xml_file(outfilename)

# Print the usage text
def usage():
    print USAGE_TEXT
    sys.exit(1)

USAGE_TEXT = """
MAEC 4.0.1 --> MAEC 4.1 XML Converter Utility

Usage: python maec_4.0.1_to_4.1.py -i <input maec 4.0.1 xml file> -o <output maec 4.1 xml file>
"""

def main():
    infilename = None
    outfilename = None

    #Get the command-line arguments
    args = sys.argv[1:]

    if len(args) < 2:
        usage()
        sys.exit(1)

    for i in range(0,len(args)):
        if args[i] == '-i':
            infilename = args[i+1]
        elif args[i] == '-o':
            outfilename = args[i+1]

    # Basic parameter checking
    if infilename and outfilename:
        update_maec(infilename, outfilename)

if __name__ == "__main__":
    main()

<commit_msg>Adds support for v4.1, renames 4.0.1 filenames<commit_after># MAEC 4.0.1 to MAEC 4.1 Converter Script
# Translates a MAEC 4.0.1 Package or Bundle into a valid MAEC 4.1 Package or Bundle

import sys
import os
import shutil
import maec
from maec.bundle.bundle import Bundle
from maec.package.package import Package

# Update the MAEC v4.0.1 file to MAEC v4.1
def update_maec(infilename, outfilename):
    # Parse the input document using the parse_xml_instance() method
    maec_objects = maec.parse_xml_instance(infilename, check_version = False)

    # Get the API Object from the parsed input
    api_object = maec_objects['api']

    # Determine if we're dealing with a Package or Bundle
    if isinstance(api_object, Package):
        # Update the Package schema_version
        api_object.schema_version = "2.1"
        for malware_subject in api_object.malware_subjects:
            for analysis in malware_subject.analyses:
                # Replace the Analysis type value of "manual" with "in-depth"
                if analysis.type and analysis.type == "manual":
                    analysis.type = "in-depth"
            # Update the schema_versions on the Bundles
            for bundle in malware_subject.findings_bundles.bundles:
                bundle.schema_version = "4.1"
    elif isinstance(api_object, Bundle):
        # Update the Bundle schema_version
        api_object.schema_version = "4.1"

    # Output the updated MAEC object to XML
    api_object.to_xml_file(outfilename)

# Print the usage text
def usage():
    print USAGE_TEXT
    sys.exit(1)

USAGE_TEXT = """
MAEC 4.0.1 --> MAEC 4.1 XML Converter Utility

Usage: python maec_4.0.1_to_4.1.py -i <input maec 4.0.1 xml file> -o <output maec 4.1 xml file>
"""

def main():
    infilename = None
    outfilename = None
    directoryname = ''
    filepath = ''

    #Get the command-line arguments
    args = sys.argv[1:]

    if len(args) < 2:
        usage()
        sys.exit(1)

    for i in range(0,len(args)):
        if args[i] == '-i':
            infilename = args[i+1]
        elif args[i] == '-o':
            outfilename = args[i+1]
        elif args[i] == '-d':
            directoryname = args[i+1]

    if directoryname != '':
        for filename in os.listdir(directoryname):
            print filename
            if '.xml' not in filename:
                pass
            elif '_report.maec-4.0.1' not in filename:
                update_maec(os.path.join(directoryname, filename), filename.rstrip('.xml') + '_cuckoobox_maec.xml')
            else:
                new_filepath = os.path.join(directoryname, filename.replace('_report.maec-4.0.1', ''))
                shutil.move(os.path.join(directoryname, filename), new_filepath)
                update_maec(new_filepath, new_filepath.rstrip('.xml') + '_cuckoobox_maec.xml')

    # Basic parameter checking
    elif infilename and outfilename:
        update_maec(infilename, outfilename)

if __name__ == "__main__":
    main()

<|endoftext|>


<commit_before>"""
Command line interface for Pushy.
Requires pushy_interpreter.py to be available.
"""

import ast
import sys

USAGE = """Usage:
  $ pushy t <script-text> [input]
  $ pushy f <script-file> [input]"""

def error(text):
    raise SystemExit(text)

try:
    import pushy_interpreter
except:
    error("Error: Could not import 'pushy_interpreter'.")

def parse_args():
    args = sys.argv[1:]

    if len(args) < 2:
        error(USAGE)

    scripttype = args[0].lower()
    if scripttype not in ('t', 'f'):
        error(USAGE)

    if scripttype == 't':
        script = args[1]
    elif scripttype == 'f':
        try:
            script = open(args[1]).read()
        except:
            error("Could not find file '{0}'.".format(args[1]))

    if len(args) == 2:
        return (script, [])

    return (script, safe_eval(args[2]))

def safe_eval(text):
    try:
        literal = ast.literal_eval(text)
    except:
        error("Could not parse input. Please provide a valid Python literal.")

    allowed_types = [int, list, tuple, str]

    if type(literal) in (tuple, list):
        for item in literal:
            if type(item) != int:
                error("Inputs lists must only contain integers. ")
        return literal

    elif type(literal) == int:
        return [literal]

    elif type(literal) == str:
        return [ord(char) for char in literal]

    else:
        error("Not a valid input type.")

if __name__ == '__main__':

    args = parse_args()
    try:
        pushy_interpreter.Script(args[0]).run(args[1])
    except KeyboardInterrupt:
        error("Program terminated (keyboard interruption).")
    except Exception as e:
        error("Python error: " + repr(e))
<commit_msg>Minor formatting<commit_after>"""
Command line interface for Pushy.
Requires pushy_interpreter.py to be available.
"""

import ast
import sys

USAGE = """Usage:
  $ pushy t <script-text> [input]
  $ pushy f <script-file> [input]"""

def error(text):
    raise SystemExit(text)

try:
    import pushy_interpreter
except:
    error("Error: Could not import 'pushy_interpreter'.")

def parse_args():
    args = sys.argv[1:]

    if len(args) < 2:
        error(USAGE)

    scripttype = args[0].lower()
    if scripttype not in ('t', 'f'):
        error(USAGE)

    if scripttype == 't':
        script = args[1]
    elif scripttype == 'f':
        try:
            script = open(args[1]).read()
        except:
            error("Could not find file '{0}'.".format(args[1]))

    if len(args) == 2:
        return (script, [])

    return (script, safe_eval(args[2]))

def safe_eval(text):
    try:
        literal = ast.literal_eval(text)
    except:
        error("Could not parse input. Please provide a valid Python literal.")

    allowed_types = [int, list, tuple, str]

    if type(literal) in (tuple, list):
        for item in literal:
            if type(item) != int:
                error("Inputs lists must only contain integers. ")
        return literal

    elif type(literal) == int:
        return [literal]

    elif type(literal) == str:
        return [ord(char) for char in literal]

    else:
        error("Not a valid input type.")

if __name__ == '__main__':

    args = parse_args()
    try:
        pushy_interpreter.Script(args[0]).run(args[1])
    except KeyboardInterrupt:
        error("\nProgram terminated (keyboard interruption).")
    except Exception as e:
        error("\nPython error: " + repr(e))
<|endoftext|>


<commit_before>#!/usr/bin/env python
#########################################################################################
#
# Apply transformations. This function is a wrapper for sct_WarpImageMultiTransform
#
# ---------------------------------------------------------------------------------------
# Copyright (c) 2014 Polytechnique Montreal <www.neuro.polymtl.ca>
# Authors: Julien Cohen-Adad, Olivier Comtois
# Modified: 2014-07-20
#
# About the license: see the file LICENSE.TXT
#########################################################################################

# TODO: display message at the end
# TODO: interpolation methods

import sys
import os
import getopt
import commands
import time
from msct_parser import Parser
import sct_utils as sct
from sct_crop_image import ImageCropper

class Param:
    def __init__(self):
        self.verbose = '1'
        self.remove_tmp_files = '1'

# PARSER
# ==========================================================================================
def get_parser():
    # parser initialisation
    parser = Parser(__file__)
    parser.usage.set_description('Apply transformations. This function is a wrapper for antsApplyTransforms (ANTs).')
    parser.add_option(name="-i",
                      type_value="file",
                      description="input image",
                      mandatory=True,
                      example="t2.nii.gz")
    parser.add_option(name="-d",
                      type_value="file",
                      description="destination image",
                      mandatory=True,
                      example="out.nii.gz")
    parser.add_option(name="-w",
                      type_value=[[','], "file"],
                      description="warping field",
                      mandatory=True,
                      example="warp1.nii.gz,warp2.nii.gz")
    parser.add_option(name="-crop",
                      type_value="int",
                      description="Crop Reference. 0 : no reference. 1 : sets background to 0. 2 : use normal background",
                      mandatory=False,
                      default_value='0',
                      example=['0','1','2'])
    parser.add_option(name="-c",
                      type_value=None,
                      description="Crop Reference. 0 : no reference. 1 : sets background to 0. 2 : use normal background",
                      mandatory=False,
                      deprecated_by='-crop')
    parser.add_option(name="-o",
                      type_value="file_output",
                      description="registered source.",
                      mandatory=False,
                      default_value='',
                      example="dest.nii.gz")
    parser.add_option(name="-x",
                      type_value="multiple_choice",
                      description="interpolation method",
                      mandatory=False,
                      default_value='spline',
                      example=['nn','linear','spline'])
    parser.add_option(name="-r",
                      type_value="multiple_choice",
                      description="""Remove temporary files.""",
                      mandatory=False,
                      default_value='1',
                      example=['0', '1'])
    parser.add_option(name="-v",
                      type_value="multiple_choice",
                      description="""Verbose.""",
                      mandatory=False,
                      default_value='0',
                      example=['0', '1', '2'])

    return parser

class Transform:
    def __init__(self, input_filename, warp, fname_dest, output_filename='', verbose=0, crop=0, interp='spline', remove_temp_files=1, debug=0):
        self.input_filename = input_filename
        if isinstance(warp, str):
            self.warp_input = list([warp])
        else:
            self.warp_input = warp
        self.fname_dest = fname_dest
        self.output_filename = output_filename
        self.interp = interp
        self.crop = crop
        self.verbose = verbose
        self.remove_temp_files = remove_temp_files
        self.debug = debug

    def apply(self):
        # Initialization
        fname_src = self.input_filename  # source image (moving)
        fname_warp_list = self.warp_input  # list of warping fields
        fname_out = self.output_filename  # output
        fname_dest = self.fname_dest  # destination image (fix)
        verbose = self.verbose
        remove_temp_files = self.remove_temp_files
        crop_reference = self.crop  # if = 1, put 0 everywhere around warping field, if = 2, real crop

        interp = sct.get_interpolation('isct_antsApplyTransforms', self.interp)

        # Parse list of warping fields
        sct.printv('\nParse list of warping fields...', verbose)
        use_inverse = []
        fname_warp_list_invert = []
        # fname_warp_list = fname_warp_list.replace(' ', '')  # remove spaces
        # fname_warp_list = fname_warp_list.split(",")  # parse with comma
        for i in range(len(fname_warp_list)):
            # Check if inverse matrix is specified with '-' at the beginning of file name
            if fname_warp_list[i].find('-') == 0:
                use_inverse.append('-i ')
                fname_warp_list[i] = fname_warp_list[i][1:]  # remove '-'
            else:
                use_inverse.append('')
            sct.printv('  Transfo #'+str(i)+': '+use_inverse[i]+fname_warp_list[i], verbose)
            fname_warp_list_invert.append(use_inverse[i]+fname_warp_list[i])

        # need to check if last warping field is an affine transfo
        isLastAffine = False
        path_fname, file_fname, ext_fname = sct.extract_fname(fname_warp_list_invert[-1])
        if ext_fname in ['.txt', '.mat']:
            isLastAffine = True

        # Check file existence
        sct.printv('\nCheck file existence...', verbose)
        sct.check_file_exist(fname_src, self.verbose)
        sct.check_file_exist(fname_dest, self.verbose)
        for i in range(len(fname_warp_list)):
            # check if file exist
            sct.check_file_exist(fname_warp_list[i], self.verbose)

        # check if destination file is 3d
        if not sct.check_if_3d(fname_dest):
            sct.printv('ERROR: Destination data must be 3d')

        # N.B. Here we take the inverse of the warp list, because sct_WarpImageMultiTransform concatenates in the reverse order
        fname_warp_list_invert.reverse()

        # Extract path, file and extension
        path_src, file_src, ext_src = sct.extract_fname(fname_src)
        path_dest, file_dest, ext_dest = sct.extract_fname(fname_dest)

        # Get output folder and file name
        if fname_out == '':
            path_out = ''  # output in user's current directory
            file_out = file_src+'_reg'
            ext_out = ext_src
            fname_out = path_out+file_out+ext_out

        # Get dimensions of data
        sct.printv('\nGet dimensions of data...', verbose)
        from msct_image import Image
        nx, ny, nz, nt, px, py, pz, pt = Image(fname_src).dim
        # nx, ny, nz, nt, px, py, pz, pt = sct.get_dimension(fname_src)
        sct.printv('  ' + str(nx) + ' x ' + str(ny) + ' x ' + str(nz)+ ' x ' + str(nt), verbose)

        # if 3d
        if nt == 1:
            # Apply transformation
            sct.printv('\nApply transformation...', verbose)
            sct.run('isct_antsApplyTransforms -d 3 -i '+fname_src+' -o '+fname_out+' -t '+' '.join(fname_warp_list_invert)+' -r '+fname_dest+interp, verbose)

        # if 4d, loop across the T dimension
        else:
            # create temporary folder
            sct.printv('\nCreate temporary folder...', verbose)
            path_tmp = sct.slash_at_the_end('tmp.'+time.strftime("%y%m%d%H%M%S"), 1)
            # sct.run('mkdir '+path_tmp, verbose)
            sct.run('mkdir '+path_tmp, verbose)

            # convert to nifti into temp folder
            sct.printv('\nCopying input data to tmp folder and convert to nii...', verbose)
            from sct_convert import convert
            convert(fname_src, path_tmp+'data.nii')
            sct.run('cp '+fname_dest+' '+path_tmp+file_dest+ext_dest)
            fname_warp_list_tmp = []
            for fname_warp in fname_warp_list:
                path_warp, file_warp, ext_warp = sct.extract_fname(fname_warp)
                sct.run('cp '+fname_warp+' '+path_tmp+file_warp+ext_warp)
                fname_warp_list_tmp.append(file_warp+ext_warp)
            fname_warp_list_invert_tmp = fname_warp_list_tmp[::-1]

            os.chdir(path_tmp)
            # split along T dimension
            sct.printv('\nSplit along T dimension...', verbose)
            from sct_image import split_data
            im_dat = Image('data.nii')
            data_split_list = split_data(im_dat, 3)
            for im in data_split_list:
                im.save()

            # apply transfo
            sct.printv('\nApply transformation to each 3D volume...', verbose)
            for it in range(nt):
                file_data_split = 'data_T'+str(it).zfill(4)+'.nii'
                file_data_split_reg = 'data_reg_T'+str(it).zfill(4)+'.nii'
                sct.run('isct_antsApplyTransforms -d 3 -i '+file_data_split+' -o '+file_data_split_reg+' -t '+' '.join(fname_warp_list_invert_tmp)+' -r '+file_dest+ext_dest+interp, verbose)

            # Merge files back
            sct.printv('\nMerge file back...', verbose)
            from sct_image import concat_data
            import glob
            path_out, name_out, ext_out = sct.extract_fname(fname_out)
            im_list = [Image(file_name) for file_name in glob.glob('data_reg_T*.nii')]
            im_out = concat_data(im_list, 3)
            im_out.setFileName(name_out+ext_out)
            im_out.save()
            os.chdir('..')
            sct.generate_output_file(path_tmp+name_out+ext_out, fname_out)
            # Delete temporary folder if specified
            if int(remove_temp_files):
                sct.printv('\nRemove temporary files...', verbose)
                sct.run('rm -rf '+path_tmp, verbose, error_exit='warning')

        # 2. crop the resulting image using dimensions from the warping field
        warping_field = fname_warp_list_invert[-1]
        # if last warping field is an affine transfo, we need to compute the space of the concatenate warping field:
        if isLastAffine:
            sct.printv('WARNING: the resulting image could have wrong apparent results. You should use an affine transformation as last transformation...',verbose,'warning')
        elif crop_reference == 1:
            ImageCropper(input_file=fname_out, output_file=fname_out, ref=warping_field, background=0).crop()
            # sct.run('sct_crop_image -i '+fname_out+' -o '+fname_out+' -ref '+warping_field+' -b 0')
        elif crop_reference == 2:
            ImageCropper(input_file=fname_out, output_file=fname_out, ref=warping_field).crop()
            # sct.run('sct_crop_image -i '+fname_out+' -o '+fname_out+' -ref '+warping_field)

        # display elapsed time
        sct.printv('\nDone! To view results, type:', verbose)
        sct.printv('fslview '+fname_dest+' '+fname_out+' &\n', verbose, 'info')

# MAIN
# ==========================================================================================
def main(args=None):

    # check user arguments
    if not args:
        args = sys.argv[1:]

    # Get parser info
    parser = get_parser()
    arguments = parser.parse(sys.argv[1:])
    input_filename = arguments["-i"]
    fname_dest = arguments["-d"]
    warp_filename = arguments["-w"]

    transform = Transform(input_filename=input_filename, fname_dest=fname_dest, warp=warp_filename)

    if "-crop" in arguments:
        transform.crop = arguments["-crop"]
    if "-o" in arguments:
        transform.output_filename = arguments["-o"]
    if "-x" in arguments:
        transform.interp = arguments["-x"]
    if "-r" in arguments:
        transform.remove_temp_files = int(arguments["-r"])
    if "-v" in arguments:
        transform.verbose = int(arguments["-v"])

    transform.apply()

# START PROGRAM
# ==========================================================================================
if __name__ == "__main__":
    # # initialize parameters
    param = Param()
    # call main function
    main()
<commit_msg>REF: -crop is a multiple choice flag<commit_after>#!/usr/bin/env python
#########################################################################################
#
# Apply transformations. This function is a wrapper for sct_WarpImageMultiTransform
#
# ---------------------------------------------------------------------------------------
# Copyright (c) 2014 Polytechnique Montreal <www.neuro.polymtl.ca>
# Authors: Julien Cohen-Adad, Olivier Comtois
# Modified: 2014-07-20
#
# About the license: see the file LICENSE.TXT
#########################################################################################

# TODO: display message at the end
# TODO: interpolation methods

import sys
import os
import getopt
import commands
import time
from msct_parser import Parser
import sct_utils as sct
from sct_crop_image import ImageCropper

class Param:
    def __init__(self):
        self.verbose = '1'
        self.remove_tmp_files = '1'

# PARSER
# ==========================================================================================
def get_parser():
    # parser initialisation
    parser = Parser(__file__)
    parser.usage.set_description('Apply transformations. This function is a wrapper for antsApplyTransforms (ANTs).')
    parser.add_option(name="-i",
                      type_value="file",
                      description="input image",
                      mandatory=True,
                      example="t2.nii.gz")
    parser.add_option(name="-d",
                      type_value="file",
                      description="destination image",
                      mandatory=True,
                      example="out.nii.gz")
    parser.add_option(name="-w",
                      type_value=[[','], "file"],
                      description="warping field",
                      mandatory=True,
                      example="warp1.nii.gz,warp2.nii.gz")
    parser.add_option(name="-crop",
                      type_value="multiple_choice",
                      description="Crop Reference. 0 : no reference. 1 : sets background to 0. 2 : use normal background",
                      mandatory=False,
                      default_value='0',
                      example=['0','1','2'])
    parser.add_option(name="-c",
                      type_value=None,
                      description="Crop Reference. 0 : no reference. 1 : sets background to 0. 2 : use normal background",
                      mandatory=False,
                      deprecated_by='-crop')
    parser.add_option(name="-o",
                      type_value="file_output",
                      description="registered source.",
                      mandatory=False,
                      default_value='',
                      example="dest.nii.gz")
    parser.add_option(name="-x",
                      type_value="multiple_choice",
                      description="interpolation method",
                      mandatory=False,
                      default_value='spline',
                      example=['nn','linear','spline'])
    parser.add_option(name="-r",
                      type_value="multiple_choice",
                      description="""Remove temporary files.""",
                      mandatory=False,
                      default_value='1',
                      example=['0', '1'])
    parser.add_option(name="-v",
                      type_value="multiple_choice",
                      description="""Verbose.""",
                      mandatory=False,
                      default_value='0',
                      example=['0', '1', '2'])

    return parser

class Transform:
    def __init__(self, input_filename, warp, fname_dest, output_filename='', verbose=0, crop=0, interp='spline', remove_temp_files=1, debug=0):
        self.input_filename = input_filename
        if isinstance(warp, str):
            self.warp_input = list([warp])
        else:
            self.warp_input = warp
        self.fname_dest = fname_dest
        self.output_filename = output_filename
        self.interp = interp
        self.crop = crop
        self.verbose = verbose
        self.remove_temp_files = remove_temp_files
        self.debug = debug

    def apply(self):
        # Initialization
        fname_src = self.input_filename  # source image (moving)
        fname_warp_list = self.warp_input  # list of warping fields
        fname_out = self.output_filename  # output
        fname_dest = self.fname_dest  # destination image (fix)
        verbose = self.verbose
        remove_temp_files = self.remove_temp_files
        crop_reference = self.crop  # if = 1, put 0 everywhere around warping field, if = 2, real crop

        interp = sct.get_interpolation('isct_antsApplyTransforms', self.interp)

        # Parse list of warping fields
        sct.printv('\nParse list of warping fields...', verbose)
        use_inverse = []
        fname_warp_list_invert = []
        # fname_warp_list = fname_warp_list.replace(' ', '')  # remove spaces
        # fname_warp_list = fname_warp_list.split(",")  # parse with comma
        for i in range(len(fname_warp_list)):
            # Check if inverse matrix is specified with '-' at the beginning of file name
            if fname_warp_list[i].find('-') == 0:
                use_inverse.append('-i ')
                fname_warp_list[i] = fname_warp_list[i][1:]  # remove '-'
            else:
                use_inverse.append('')
            sct.printv('  Transfo #'+str(i)+': '+use_inverse[i]+fname_warp_list[i], verbose)
            fname_warp_list_invert.append(use_inverse[i]+fname_warp_list[i])

        # need to check if last warping field is an affine transfo
        isLastAffine = False
        path_fname, file_fname, ext_fname = sct.extract_fname(fname_warp_list_invert[-1])
        if ext_fname in ['.txt', '.mat']:
            isLastAffine = True

        # Check file existence
        sct.printv('\nCheck file existence...', verbose)
        sct.check_file_exist(fname_src, self.verbose)
        sct.check_file_exist(fname_dest, self.verbose)
        for i in range(len(fname_warp_list)):
            # check if file exist
            sct.check_file_exist(fname_warp_list[i], self.verbose)

        # check if destination file is 3d
        if not sct.check_if_3d(fname_dest):
            sct.printv('ERROR: Destination data must be 3d')

        # N.B. Here we take the inverse of the warp list, because sct_WarpImageMultiTransform concatenates in the reverse order
        fname_warp_list_invert.reverse()

        # Extract path, file and extension
        path_src, file_src, ext_src = sct.extract_fname(fname_src)
        path_dest, file_dest, ext_dest = sct.extract_fname(fname_dest)

        # Get output folder and file name
        if fname_out == '':
            path_out = ''  # output in user's current directory
            file_out = file_src+'_reg'
            ext_out = ext_src
            fname_out = path_out+file_out+ext_out

        # Get dimensions of data
        sct.printv('\nGet dimensions of data...', verbose)
        from msct_image import Image
        nx, ny, nz, nt, px, py, pz, pt = Image(fname_src).dim
        # nx, ny, nz, nt, px, py, pz, pt = sct.get_dimension(fname_src)
        sct.printv('  ' + str(nx) + ' x ' + str(ny) + ' x ' + str(nz)+ ' x ' + str(nt), verbose)

        # if 3d
        if nt == 1:
            # Apply transformation
            sct.printv('\nApply transformation...', verbose)
            sct.run('isct_antsApplyTransforms -d 3 -i '+fname_src+' -o '+fname_out+' -t '+' '.join(fname_warp_list_invert)+' -r '+fname_dest+interp, verbose)

        # if 4d, loop across the T dimension
        else:
            # create temporary folder
            sct.printv('\nCreate temporary folder...', verbose)
            path_tmp = sct.slash_at_the_end('tmp.'+time.strftime("%y%m%d%H%M%S"), 1)
            # sct.run('mkdir '+path_tmp, verbose)
            sct.run('mkdir '+path_tmp, verbose)

            # convert to nifti into temp folder
            sct.printv('\nCopying input data to tmp folder and convert to nii...', verbose)
            from sct_convert import convert
            convert(fname_src, path_tmp+'data.nii')
            sct.run('cp '+fname_dest+' '+path_tmp+file_dest+ext_dest)
            fname_warp_list_tmp = []
            for fname_warp in fname_warp_list:
                path_warp, file_warp, ext_warp = sct.extract_fname(fname_warp)
                sct.run('cp '+fname_warp+' '+path_tmp+file_warp+ext_warp)
                fname_warp_list_tmp.append(file_warp+ext_warp)
            fname_warp_list_invert_tmp = fname_warp_list_tmp[::-1]

            os.chdir(path_tmp)
            # split along T dimension
            sct.printv('\nSplit along T dimension...', verbose)
            from sct_image import split_data
            im_dat = Image('data.nii')
            data_split_list = split_data(im_dat, 3)
            for im in data_split_list:
                im.save()

            # apply transfo
            sct.printv('\nApply transformation to each 3D volume...', verbose)
            for it in range(nt):
                file_data_split = 'data_T'+str(it).zfill(4)+'.nii'
                file_data_split_reg = 'data_reg_T'+str(it).zfill(4)+'.nii'
                sct.run('isct_antsApplyTransforms -d 3 -i '+file_data_split+' -o '+file_data_split_reg+' -t '+' '.join(fname_warp_list_invert_tmp)+' -r '+file_dest+ext_dest+interp, verbose)

            # Merge files back
            sct.printv('\nMerge file back...', verbose)
            from sct_image import concat_data
            import glob
            path_out, name_out, ext_out = sct.extract_fname(fname_out)
            im_list = [Image(file_name) for file_name in glob.glob('data_reg_T*.nii')]
            im_out = concat_data(im_list, 3)
            im_out.setFileName(name_out+ext_out)
            im_out.save()
            os.chdir('..')
            sct.generate_output_file(path_tmp+name_out+ext_out, fname_out)
            # Delete temporary folder if specified
            if int(remove_temp_files):
                sct.printv('\nRemove temporary files...', verbose)
                sct.run('rm -rf '+path_tmp, verbose, error_exit='warning')

        # 2. crop the resulting image using dimensions from the warping field
        warping_field = fname_warp_list_invert[-1]
        # if last warping field is an affine transfo, we need to compute the space of the concatenate warping field:
        if isLastAffine:
            sct.printv('WARNING: the resulting image could have wrong apparent results. You should use an affine transformation as last transformation...',verbose,'warning')
        elif crop_reference == 1:
            ImageCropper(input_file=fname_out, output_file=fname_out, ref=warping_field, background=0).crop()
            # sct.run('sct_crop_image -i '+fname_out+' -o '+fname_out+' -ref '+warping_field+' -b 0')
        elif crop_reference == 2:
            ImageCropper(input_file=fname_out, output_file=fname_out, ref=warping_field).crop()
            # sct.run('sct_crop_image -i '+fname_out+' -o '+fname_out+' -ref '+warping_field)

        # display elapsed time
        sct.printv('\nDone! To view results, type:', verbose)
        sct.printv('fslview '+fname_dest+' '+fname_out+' &\n', verbose, 'info')

# MAIN
# ==========================================================================================
def main(args=None):

    # check user arguments
    if not args:
        args = sys.argv[1:]

    # Get parser info
    parser = get_parser()
    arguments = parser.parse(sys.argv[1:])
    input_filename = arguments["-i"]
    fname_dest = arguments["-d"]
    warp_filename = arguments["-w"]

    transform = Transform(input_filename=input_filename, fname_dest=fname_dest, warp=warp_filename)

    if "-crop" in arguments:
        transform.crop = arguments["-crop"]
    if "-o" in arguments:
        transform.output_filename = arguments["-o"]
    if "-x" in arguments:
        transform.interp = arguments["-x"]
    if "-r" in arguments:
        transform.remove_temp_files = int(arguments["-r"])
    if "-v" in arguments:
        transform.verbose = int(arguments["-v"])

    transform.apply()

# START PROGRAM
# ==========================================================================================
if __name__ == "__main__":
    # # initialize parameters
    param = Param()
    # call main function
    main()
<|endoftext|>


<commit_before>"""
Pytest test running.

This module implements the ``test()`` function for NumPy modules. The usual
boiler plate for doing that is to put the following in the module
``__init__.py`` file::

    from numpy._pytesttester import PytestTester
    test = PytestTester(__name__)
    del PytestTester

Warnings filtering and other runtime settings should be dealt with in the
``pytest.ini`` file in the numpy repo root. The behavior of the test depends on
whether or not that file is found as follows:

* ``pytest.ini`` is present (develop mode)
    All warnings except those explicitly filtered out are raised as error.
* ``pytest.ini`` is absent (release mode)
    DeprecationWarnings and PendingDeprecationWarnings are ignored, other
    warnings are passed through.

In practice, tests run from the numpy repo are run in develop mode. That
includes the standard ``python runtests.py`` invocation.

This module is imported by every numpy subpackage, so lies at the top level to
simplify circular import issues. For the same reason, it contains no numpy
imports at module scope, instead importing numpy within function calls.
"""
import sys
import os

__all__ = ['PytestTester']

def _show_numpy_info():
    import numpy as np

    print("NumPy version %s" % np.__version__)
    relaxed_strides = np.ones((10, 1), order="C").flags.f_contiguous
    print("NumPy relaxed strides checking option:", relaxed_strides)
    info = np.lib.utils._opt_info()
    print("NumPy CPU features: ", (info if info else 'nothing enabled'))

class PytestTester:
    """
    Pytest test runner.

    A test function is typically added to a package's __init__.py like so::

      from numpy._pytesttester import PytestTester
      test = PytestTester(__name__).test
      del PytestTester

    Calling this test function finds and runs all tests associated with the
    module and all its sub-modules.

    Attributes
    ----------
    module_name : str
        Full path to the package to test.

    Parameters
    ----------
    module_name : module name
        The name of the module to test.

    Notes
    -----
    Unlike the previous ``nose``-based implementation, this class is not
    publicly exposed as it performs some ``numpy``-specific warning
    suppression.

    """
    def __init__(self, module_name):
        self.module_name = module_name

    def __call__(self, label='fast', verbose=1, extra_argv=None,
                 doctests=False, coverage=False, durations=-1, tests=None):
        """
        Run tests for module using pytest.

        Parameters
        ----------
        label : {'fast', 'full'}, optional
            Identifies the tests to run. When set to 'fast', tests decorated
            with `pytest.mark.slow` are skipped, when 'full', the slow marker
            is ignored.
        verbose : int, optional
            Verbosity value for test outputs, in the range 1-3. Default is 1.
        extra_argv : list, optional
            List with any extra arguments to pass to pytests.
        doctests : bool, optional
            .. note:: Not supported
        coverage : bool, optional
            If True, report coverage of NumPy code. Default is False.
            Requires installation of (pip) pytest-cov.
        durations : int, optional
            If < 0, do nothing, If 0, report time of all tests, if > 0,
            report the time of the slowest `timer` tests. Default is -1.
        tests : test or list of tests
            Tests to be executed with pytest '--pyargs'

        Returns
        -------
        result : bool
            Return True on success, false otherwise.

        Notes
        -----
        Each NumPy module exposes `test` in its namespace to run all tests for
        it. For example, to run all tests for numpy.lib:

        >>> np.lib.test() #doctest: +SKIP

        Examples
        --------
        >>> result = np.lib.test() #doctest: +SKIP
        ...
        1023 passed, 2 skipped, 6 deselected, 1 xfailed in 10.39 seconds
        >>> result
        True

        """
        import pytest
        import warnings

        module = sys.modules[self.module_name]
        module_path = os.path.abspath(module.__path__[0])

        # setup the pytest arguments
        pytest_args = ["-l"]

        # offset verbosity. The "-q" cancels a "-v".
        pytest_args += ["-q"]

        with warnings.catch_warnings():
            warnings.simplefilter("always")
            # Filter out distutils cpu warnings (could be localized to
            # distutils tests). ASV has problems with top level import,
            # so fetch module for suppression here.
            from numpy.distutils import cpuinfo

        with warnings.catch_warnings(record=True):
            # Ignore the warning from importing the array_api submodule. This
            # warning is done on import, so it would break pytest collection,
            # but importing it early here prevents the warning from being
            # issued when it imported again.
            import numpy.array_api

        # Filter out annoying import messages. Want these in both develop and
        # release mode.
        pytest_args += [
            "-W ignore:Not importing directory",
            "-W ignore:numpy.dtype size changed",
            "-W ignore:numpy.ufunc size changed",
            "-W ignore::UserWarning:cpuinfo",
            ]

        # When testing matrices, ignore their PendingDeprecationWarnings
        pytest_args += [
            "-W ignore:the matrix subclass is not",
            "-W ignore:Importing from numpy.matlib is",
            ]

        if doctests:
            raise ValueError("Doctests not supported")

        if extra_argv:
            pytest_args += list(extra_argv)

        if verbose > 1:
            pytest_args += ["-" + "v"*(verbose - 1)]

        if coverage:
            pytest_args += ["--cov=" + module_path]

        if label == "fast":
            # not importing at the top level to avoid circular import of module
            from numpy.testing import IS_PYPY
            if IS_PYPY:
                pytest_args += ["-m", "not slow and not slow_pypy"]
            else:
                pytest_args += ["-m", "not slow"]

        elif label != "full":
            pytest_args += ["-m", label]

        if durations >= 0:
            pytest_args += ["--durations=%s" % durations]

        if tests is None:
            tests = [self.module_name]

        pytest_args += ["--pyargs"] + list(tests)

        # run tests.
        _show_numpy_info()

        try:
            code = pytest.main(pytest_args)
        except SystemExit as exc:
            code = exc.code

        return code == 0
<commit_msg>BUG: Adding the default pytest doctest runner instead of the ValueError to the np testrunner<commit_after>"""
Pytest test running.

This module implements the ``test()`` function for NumPy modules. The usual
boiler plate for doing that is to put the following in the module
``__init__.py`` file::

    from numpy._pytesttester import PytestTester
    test = PytestTester(__name__)
    del PytestTester

Warnings filtering and other runtime settings should be dealt with in the
``pytest.ini`` file in the numpy repo root. The behavior of the test depends on
whether or not that file is found as follows:

* ``pytest.ini`` is present (develop mode)
    All warnings except those explicitly filtered out are raised as error.
* ``pytest.ini`` is absent (release mode)
    DeprecationWarnings and PendingDeprecationWarnings are ignored, other
    warnings are passed through.

In practice, tests run from the numpy repo are run in develop mode. That
includes the standard ``python runtests.py`` invocation.

This module is imported by every numpy subpackage, so lies at the top level to
simplify circular import issues. For the same reason, it contains no numpy
imports at module scope, instead importing numpy within function calls.
"""
import sys
import os

__all__ = ['PytestTester']

def _show_numpy_info():
    import numpy as np

    print("NumPy version %s" % np.__version__)
    relaxed_strides = np.ones((10, 1), order="C").flags.f_contiguous
    print("NumPy relaxed strides checking option:", relaxed_strides)
    info = np.lib.utils._opt_info()
    print("NumPy CPU features: ", (info if info else 'nothing enabled'))

class PytestTester:
    """
    Pytest test runner.

    A test function is typically added to a package's __init__.py like so::

      from numpy._pytesttester import PytestTester
      test = PytestTester(__name__).test
      del PytestTester

    Calling this test function finds and runs all tests associated with the
    module and all its sub-modules.

    Attributes
    ----------
    module_name : str
        Full path to the package to test.

    Parameters
    ----------
    module_name : module name
        The name of the module to test.

    Notes
    -----
    Unlike the previous ``nose``-based implementation, this class is not
    publicly exposed as it performs some ``numpy``-specific warning
    suppression.

    """
    def __init__(self, module_name):
        self.module_name = module_name

    def __call__(self, label='fast', verbose=1, extra_argv=None,
                 doctests=False, coverage=False, durations=-1, tests=None):
        """
        Run tests for module using pytest.

        Parameters
        ----------
        label : {'fast', 'full'}, optional
            Identifies the tests to run. When set to 'fast', tests decorated
            with `pytest.mark.slow` are skipped, when 'full', the slow marker
            is ignored.
        verbose : int, optional
            Verbosity value for test outputs, in the range 1-3. Default is 1.
        extra_argv : list, optional
            List with any extra arguments to pass to pytests.
        doctests : bool, optional
            .. note:: Not supported
        coverage : bool, optional
            If True, report coverage of NumPy code. Default is False.
            Requires installation of (pip) pytest-cov.
        durations : int, optional
            If < 0, do nothing, If 0, report time of all tests, if > 0,
            report the time of the slowest `timer` tests. Default is -1.
        tests : test or list of tests
            Tests to be executed with pytest '--pyargs'

        Returns
        -------
        result : bool
            Return True on success, false otherwise.

        Notes
        -----
        Each NumPy module exposes `test` in its namespace to run all tests for
        it. For example, to run all tests for numpy.lib:

        >>> np.lib.test() #doctest: +SKIP

        Examples
        --------
        >>> result = np.lib.test() #doctest: +SKIP
        ...
        1023 passed, 2 skipped, 6 deselected, 1 xfailed in 10.39 seconds
        >>> result
        True

        """
        import pytest
        import warnings

        module = sys.modules[self.module_name]
        module_path = os.path.abspath(module.__path__[0])

        # setup the pytest arguments
        pytest_args = ["-l"]

        # offset verbosity. The "-q" cancels a "-v".
        pytest_args += ["-q"]

        with warnings.catch_warnings():
            warnings.simplefilter("always")
            # Filter out distutils cpu warnings (could be localized to
            # distutils tests). ASV has problems with top level import,
            # so fetch module for suppression here.
            from numpy.distutils import cpuinfo

        with warnings.catch_warnings(record=True):
            # Ignore the warning from importing the array_api submodule. This
            # warning is done on import, so it would break pytest collection,
            # but importing it early here prevents the warning from being
            # issued when it imported again.
            import numpy.array_api

        # Filter out annoying import messages. Want these in both develop and
        # release mode.
        pytest_args += [
            "-W ignore:Not importing directory",
            "-W ignore:numpy.dtype size changed",
            "-W ignore:numpy.ufunc size changed",
            "-W ignore::UserWarning:cpuinfo",
            ]

        # When testing matrices, ignore their PendingDeprecationWarnings
        pytest_args += [
            "-W ignore:the matrix subclass is not",
            "-W ignore:Importing from numpy.matlib is",
            ]

        if doctests:
            pytest_args += ["--doctest-modules"]

        if extra_argv:
            pytest_args += list(extra_argv)

        if verbose > 1:
            pytest_args += ["-" + "v"*(verbose - 1)]

        if coverage:
            pytest_args += ["--cov=" + module_path]

        if label == "fast":
            # not importing at the top level to avoid circular import of module
            from numpy.testing import IS_PYPY
            if IS_PYPY:
                pytest_args += ["-m", "not slow and not slow_pypy"]
            else:
                pytest_args += ["-m", "not slow"]

        elif label != "full":
            pytest_args += ["-m", label]

        if durations >= 0:
            pytest_args += ["--durations=%s" % durations]

        if tests is None:
            tests = [self.module_name]

        pytest_args += ["--pyargs"] + list(tests)

        # run tests.
        _show_numpy_info()

        try:
            code = pytest.main(pytest_args)
        except SystemExit as exc:
            code = exc.code

        return code == 0
<|endoftext|>


<commit_before># Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path
from typing import Any, Dict, List
from pandas import DataFrame, concat
from lib.concurrent import thread_map
from lib.data_source import DataSource

def _normalize_column_name(column: str) -> str:
    column = column.lower()
    column = column.replace("symptom:", "")
    column = column.replace("–", "_")
    column = column.replace(" ", "_")
    column = column.replace("-", "_")
    column = column.replace("–", "_")
    column = column.replace("'", "")
    return column

def _rename_columns(data: DataFrame) -> DataFrame:
    column_adapter = {
        "date": "date",
        "country_region_code": "country_code",
        "sub_region_1": "subregion1_name",
        "sub_region_2": "subregion2_name",
        "sub_region_1_code": "subregion1_code",
        "sub_region_2_code": "subregion2_code",
    }
    data = data.rename(columns=column_adapter)
    data.columns = [
        col if col in column_adapter.values() else "search_trends_" + _normalize_column_name(col)
        for col in data.columns
    ]

    return data

def _parse_subregion_code(code: str) -> str:
    return code.split("-")[-1]

def _process_chunk(data: DataFrame) -> DataFrame:
    data = _rename_columns(data)

    # Keep only regions which are not metro areas
    if "metro_area" in data.columns:
        data = data[data["metro_area"].isna()]

    # We can derive the keys directly for all data
    data["key"] = None
    subregion1_isna_mask = data.subregion1_code.isna()
    subregion2_isna_mask = data.subregion2_code.isna()

    # Country level has null subregions 1 and 2
    country_level_mask = subregion1_isna_mask & subregion2_isna_mask
    data.loc[country_level_mask, "key"] = data.loc[country_level_mask, "country_code"]

    state_level_mask = ~subregion1_isna_mask & subregion2_isna_mask
    data.loc[state_level_mask, "key"] = (
        data.loc[state_level_mask, "country_code"]
        + "_"
        + data.loc[state_level_mask, "subregion1_code"].apply(_parse_subregion_code)
    )

    county_level_mask = ~subregion1_isna_mask & ~subregion2_isna_mask
    data.loc[county_level_mask, "key"] = (
        data.loc[county_level_mask, "country_code"]
        + "_"
        + data.loc[county_level_mask, "subregion1_code"].apply(_parse_subregion_code)
        + "_"
        + data.loc[county_level_mask, "subregion2_code"].apply(_parse_subregion_code)
    )

    return data

class GoogleSearchTrendsDataSource(DataSource):
    def fetch(
        self,
        output_folder: Path,
        cache: Dict[str, str],
        fetch_opts: List[Dict[str, Any]],
        skip_existing: bool = False,
    ) -> Dict[str, str]:
        base_opts = dict(fetch_opts[0])
        url_tpl: str = base_opts.pop("url")

        # Get the files to process from the `parse` config options
        country_codes: List[str] = self.config["parse"]["country_codes"]
        subregion_names: List[str] = self.config["parse"].get("subregion_names")

        # Generate the URLs for the files to download from the template
        url_list = {}
        for year in ("2020", "2021"):
            for country_code in country_codes:
                if not subregion_names:
                    key = country_code
                    url_list[f"{key}_{year}"] = url_tpl.format(
                        level="country", region_code=key, year=year
                    )
                else:
                    for subregion_name in subregion_names:
                        key = f"{country_code}_{subregion_name.replace(' ', '_')}"
                        url_list[f"{key}_{year}"] = url_tpl.format(
                            level="sub_region_1", region_code=key, year=year
                        )

        # Replace the fetch options with our own processed options
        fetch_opts = [{"name": key, "url": url, **base_opts} for key, url in url_list.items()]
        return super().fetch(output_folder, cache, fetch_opts, skip_existing=skip_existing)

    def parse_dataframes(
        self, dataframes: Dict[str, DataFrame], aux: Dict[str, DataFrame], **parse_opts
    ) -> DataFrame:
        google_keys = aux["google_key_map"].set_index("google_location_key")["key"].to_dict()
        data = concat(thread_map(_process_chunk, dataframes.values(), total=len(dataframes)))
        data["key"] = data["key"].apply(lambda x: google_keys.get(x, x))
        return data.dropna(subset=["key"])
<commit_msg>Add 2022 data for search trends<commit_after># Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path
from typing import Any, Dict, List
from pandas import DataFrame, concat
from lib.concurrent import thread_map
from lib.data_source import DataSource

def _normalize_column_name(column: str) -> str:
    column = column.lower()
    column = column.replace("symptom:", "")
    column = column.replace("–", "_")
    column = column.replace(" ", "_")
    column = column.replace("-", "_")
    column = column.replace("–", "_")
    column = column.replace("'", "")
    return column

def _rename_columns(data: DataFrame) -> DataFrame:
    column_adapter = {
        "date": "date",
        "country_region_code": "country_code",
        "sub_region_1": "subregion1_name",
        "sub_region_2": "subregion2_name",
        "sub_region_1_code": "subregion1_code",
        "sub_region_2_code": "subregion2_code",
    }
    data = data.rename(columns=column_adapter)
    data.columns = [
        col if col in column_adapter.values() else "search_trends_" + _normalize_column_name(col)
        for col in data.columns
    ]

    return data

def _parse_subregion_code(code: str) -> str:
    return code.split("-")[-1]

def _process_chunk(data: DataFrame) -> DataFrame:
    data = _rename_columns(data)

    # Keep only regions which are not metro areas
    if "metro_area" in data.columns:
        data = data[data["metro_area"].isna()]

    # We can derive the keys directly for all data
    data["key"] = None
    subregion1_isna_mask = data.subregion1_code.isna()
    subregion2_isna_mask = data.subregion2_code.isna()

    # Country level has null subregions 1 and 2
    country_level_mask = subregion1_isna_mask & subregion2_isna_mask
    data.loc[country_level_mask, "key"] = data.loc[country_level_mask, "country_code"]

    state_level_mask = ~subregion1_isna_mask & subregion2_isna_mask
    data.loc[state_level_mask, "key"] = (
        data.loc[state_level_mask, "country_code"]
        + "_"
        + data.loc[state_level_mask, "subregion1_code"].apply(_parse_subregion_code)
    )

    county_level_mask = ~subregion1_isna_mask & ~subregion2_isna_mask
    data.loc[county_level_mask, "key"] = (
        data.loc[county_level_mask, "country_code"]
        + "_"
        + data.loc[county_level_mask, "subregion1_code"].apply(_parse_subregion_code)
        + "_"
        + data.loc[county_level_mask, "subregion2_code"].apply(_parse_subregion_code)
    )

    return data

class GoogleSearchTrendsDataSource(DataSource):
    def fetch(
        self,
        output_folder: Path,
        cache: Dict[str, str],
        fetch_opts: List[Dict[str, Any]],
        skip_existing: bool = False,
    ) -> Dict[str, str]:
        base_opts = dict(fetch_opts[0])
        url_tpl: str = base_opts.pop("url")

        # Get the files to process from the `parse` config options
        country_codes: List[str] = self.config["parse"]["country_codes"]
        subregion_names: List[str] = self.config["parse"].get("subregion_names")

        # Generate the URLs for the files to download from the template
        url_list = {}
        for year in ("2020", "2021", "2022"):
            for country_code in country_codes:
                if not subregion_names:
                    key = country_code
                    url_list[f"{key}_{year}"] = url_tpl.format(
                        level="country", region_code=key, year=year
                    )
                else:
                    for subregion_name in subregion_names:
                        key = f"{country_code}_{subregion_name.replace(' ', '_')}"
                        url_list[f"{key}_{year}"] = url_tpl.format(
                            level="sub_region_1", region_code=key, year=year
                        )

        # Replace the fetch options with our own processed options
        fetch_opts = [{"name": key, "url": url, **base_opts} for key, url in url_list.items()]
        return super().fetch(output_folder, cache, fetch_opts, skip_existing=skip_existing)

    def parse_dataframes(
        self, dataframes: Dict[str, DataFrame], aux: Dict[str, DataFrame], **parse_opts
    ) -> DataFrame:
        google_keys = aux["google_key_map"].set_index("google_location_key")["key"].to_dict()
        data = concat(thread_map(_process_chunk, dataframes.values(), total=len(dataframes)))
        data["key"] = data["key"].apply(lambda x: google_keys.get(x, x))
        return data.dropna(subset=["key"])
<|endoftext|>