<commit_before>from collections import defaultdict
import shutil

from defaults import *
from vcf_utils import *
from sv_interval import SVInterval, get_gaps_file, interval_overlaps_interval_list, merge_intervals
from pindel_reader import PindelReader
from breakdancer_reader import BreakDancerReader
from breakseq_reader import BreakSeqReader
from cnvnator_reader import CNVnatorReader
from generate_sv_intervals import parallel_generate_sc_intervals
from run_spades import run_spades_parallel
from run_age import run_age_parallel
from generate_final_vcf import convert_metasv_bed_to_vcf
from fasta_utils import get_contigs
from genotype import parallel_genotype_intervals

FORMAT = '%(levelname)s %(asctime)-15s %(name)-20s %(message)s'
logging.basicConfig(level=logging.INFO, format=FORMAT)
logger = logging.getLogger(__name__)

def create_dirs(dirlist):
    for dirname in dirlist:
        if not os.path.isdir(dirname):
            logger.info("Creating directory %s" % (dirname))
            os.makedirs(dirname)

def run_metasv(args):
    # Check if there is work to do
    if not (
                                args.pindel_vcf + args.breakdancer_vcf + args.breakseq_vcf + args.cnvnator_vcf + args.pindel_native + args.breakdancer_native + args.breakseq_native + args.cnvnator_native):
        logger.warning("Nothing to merge since no SV file specified")

    # Simple check for arguments
    if not args.disable_assembly:
        if not args.spades:
            logger.error("Spades executable not specified")
            return os.EX_USAGE

        if not args.age:
            logger.error("AGE executable not specified")
            return os.EX_USAGE

    # Create the directories for working
    bedtools_tmpdir = os.path.join(args.workdir, "bedtools")
    create_dirs([args.workdir, args.outdir, bedtools_tmpdir])

    # Reference handling
    if not os.path.isfile(args.reference + ".fai"):
        logger.error("Reference file %s is not indexed" % (args.reference))
        return 1

    fasta_handle = pysam.Fastafile(args.reference) if os.path.isfile(args.reference) else None
    contigs = get_contigs(args.reference)
    include_intervals = sorted(
        [SVInterval(contig.name, 0, contig.length, contig.name, "include", length=contig.length) for contig in contigs])

    # Generate the list of contigs to process
    contig_whitelist = set(args.chromosomes) if args.chromosomes else set([contig.name for contig in contigs])
    if args.keep_standard_contigs:
        contig_whitelist &= set(
            [str(i) for i in xrange(1, 23)] + ["chr%d" % (i) for i in xrange(1, 23)] + ["X", "Y", "MT", "chrX", "chrY",
                                                                                "chrM"])
    logger.info("Only SVs on the following contigs will be reported: %s" % (sorted(list(contig_whitelist))))

    # Load the intervals from different files
    vcf_name_list = [("CNVnator", args.cnvnator_vcf), ("Pindel", args.pindel_vcf),
                     ("BreakDancer", args.breakdancer_vcf),
                     ("BreakSeq", args.breakseq_vcf), ("HaplotypeCaller", args.gatk_vcf)]
    native_name_list = [("CNVnator", args.cnvnator_native, CNVnatorReader),
                        ("Pindel", args.pindel_native, PindelReader),
                        ("BreakSeq", args.breakseq_native, BreakSeqReader),
                        ("BreakDancer", args.breakdancer_native, BreakDancerReader)]

    tools = []
    intervals = {}
    sv_types = set()

    gap_intervals = []
    if args.filter_gaps:
        gaps = args.gaps if args.gaps else get_gaps_file(contig_whitelist)
        gap_intervals = sorted(load_gap_intervals(gaps))

    # Handles native input
    logger.info("Load native files")
    for toolname, nativename, svReader in native_name_list:
        # If no native file is given, ignore the tool
        if not nativename: continue

        tools.append(toolname)
        intervals[toolname] = defaultdict(list)

        for native_file in nativename:
            for record in svReader(native_file, svs_to_report=args.svs_to_report):
                interval = record.to_sv_interval()

                if not interval:
                    # This is the case for SVs we want to skip
                    continue
                if not interval_overlaps_interval_list(interval, gap_intervals) and interval.chrom in contig_whitelist:

                    # Check length
                    if interval.length < args.minsvlen:
                        continue

                    # Set wiggle
                    interval.wiggle = max(args.inswiggle if interval.sv_type == "INS" else 0, args.wiggle)

                    intervals[toolname][interval.sv_type].append(interval)

        sv_types |= set(intervals[toolname].keys())

    # Handles the VCF input cases, we will just deal with these cases
    logger.info("Load VCF files")
    for toolname, vcfname in vcf_name_list:
        # If no VCF is given, ignore the tool
        if not vcfname:
            continue

        tools.append(toolname)
        intervals[toolname] = {}

        vcf_list = []
        for vcffile in vcfname:
            if os.path.isdir(vcffile):
                logger.info("Will load from per-chromosome VCFs from directory %s for tool %s" % (vcffile, toolname))
                vcf_list += [os.path.join(vcffile, "%s.vcf.gz" % contig.name) for contig in contigs if
                             (not contig_whitelist or contig.name in contig_whitelist)]
            else:
                vcf_list.append(vcffile)

        for vcffile in vcf_list:
            load_intervals(vcffile, intervals[toolname], gap_intervals, include_intervals, toolname, contig_whitelist,
                           minsvlen=args.minsvlen, wiggle=args.wiggle, inswiggle=args.inswiggle,
                           svs_to_report=args.svs_to_report)
        sv_types |= set(intervals[toolname].keys())

    logger.info("SV types are %s" % (str(sv_types)))
    tool_merged_intervals = {}
    final_intervals = []

    # This will just output per-tool VCFs, no intra-tool merging is done yet
    if args.enable_per_tool_output:
        logger.info("Output per-tool VCFs")
        for toolname in intervals:
            tool_out = os.path.join(args.outdir, "%s.vcf" % (toolname.lower()))

            logger.info("Outputting single tool VCF for %s" % (str(toolname)))
            vcf_template_reader = vcf.Reader(open(os.path.join(mydir, "resources/template.vcf"), "r"))
            vcf_template_reader.samples = [args.sample]

            intervals_tool = []
            tool_out_fd = open(tool_out, "w")
            vcf_writer = vcf.Writer(tool_out_fd, vcf_template_reader)
            chr_intervals_tool = {contig.name: [] for contig in contigs}
            for sv_type in sv_types:
                if sv_type in intervals[toolname]:
                    intervals_tool.extend([copy.deepcopy(interval) for interval in intervals[toolname][sv_type]])
            for interval in intervals_tool:
                # Marghoob says that this is just to fill-in some metadata
                interval.do_validation(args.overlap_ratio)

                interval.fix_pos()
                chr_intervals_tool[interval.chrom].append(interval)

            for contig in contigs:
                chr_intervals_tool[contig.name].sort()
                for interval in chr_intervals_tool[contig.name]:
                    vcf_record = interval.to_vcf_record(fasta_handle, args.sample)
                    if vcf_record is not None:
                        vcf_writer.write_record(vcf_record)
            tool_out_fd.close()
            vcf_writer.close()
            logger.info("Indexing single tool VCF for %s" % (str(toolname)))
            pysam.tabix_index(tool_out, force=True, preset="vcf")

    # Do merging here
    logger.info("Do merging")
    for sv_type in sv_types:
        logger.info("Processing SVs of type %s" % sv_type)
        tool_merged_intervals[sv_type] = []

        # Do the intra-tool merging
        logger.info("Intra-tool Merging SVs of type %s" % sv_type)
        for tool in tools:
            logger.debug("Is %s in tool keys? %s" % (sv_type, str(intervals[tool].keys())))
            if sv_type not in intervals[tool]:
                logger.debug("%s not in tool %s" % (sv_type, tool))
                continue
            logger.info("First level merging for %s for tool %s" % (sv_type, tool))
            tool_merged_intervals[sv_type] += merge_intervals(intervals[tool][sv_type])

        # Do the inter-tool merging
        logger.info("Inter-tool Merging SVs of type %s" % sv_type)
        merged_intervals = merge_intervals(tool_merged_intervals[sv_type])

        # Intervals which overlap well with merged_intervals
        intervals1 = []
        # Intervals which do not overlap well with merged_intervals.
        # Used to filter out small intervals which got merged with large intervals
        intervals2 = []

        logger.info("Checking overlaps SVs of type %s" % sv_type)
        for interval in tool_merged_intervals[sv_type]:
            if interval_overlaps_interval_list(interval, merged_intervals, args.overlap_ratio, args.overlap_ratio):
                intervals2.append(interval)
            else:
                intervals1.append(interval)
        final_intervals.extend(merge_intervals(intervals1) + merge_intervals(intervals2))

    final_chr_intervals = {contig.name: [] for contig in contigs}
    for interval in final_intervals:
        interval.do_validation(args.overlap_ratio)
        interval.fix_pos()
        if args.minsvlen <= interval.length <= args.maxsvlen:
            final_chr_intervals[interval.chrom].append(interval)

    # This is the merged VCF without assembly, ok for deletions at this point
    logger.info("Output merged VCF without assembly ")
    vcf_template_reader = vcf.Reader(open(os.path.join(mydir, "resources/template.vcf"), "r"))
    vcf_template_reader.samples = [args.sample]
    preasm_vcf = os.path.join(args.workdir, "pre_asm.vcf")
    vcf_fd = open(preasm_vcf, "w")
    vcf_writer = vcf.Writer(vcf_fd, vcf_template_reader)

    final_stats = {}

    bed_intervals = []
    for contig in contigs:
        final_chr_intervals[contig.name].sort()
        for interval in final_chr_intervals[contig.name]:
            vcf_record = interval.to_vcf_record(fasta_handle)
            if vcf_record is not None:
                key = (interval.sv_type, "PASS" if interval.is_validated else "LowQual",
                       "PRECISE" if interval.is_precise else "IMPRECISE", tuple(sorted(list(interval.sources))))
                if key not in final_stats:
                    final_stats[key] = 0
                final_stats[key] += 1
                vcf_writer.write_record(vcf_record)
            bed_interval = interval.to_bed_interval(args.sample)
            if bed_interval is not None:
                bed_intervals.append(bed_interval)
    vcf_fd.close()
    vcf_writer.close()

    # Also save a BED file representation of the merged variants without assembly
    merged_bed = None
    if bed_intervals:
        merged_bed = os.path.join(args.workdir, "metasv.bed")
        pybedtools.BedTool(bed_intervals).saveas(merged_bed)

    for key in sorted(final_stats.keys()):
        logger.info(str(key) + ":" + str(final_stats[key]))

    final_vcf = os.path.join(args.outdir, "variants.vcf")

    # Run assembly here
    if not args.disable_assembly:
        logger.info("Running assembly")

        spades_tmpdir = os.path.join(args.workdir, "spades")
        age_tmpdir = os.path.join(args.workdir, "age")

        create_dirs([spades_tmpdir, age_tmpdir])

        assembly_bed = merged_bed

        # this does the improved assembly location finder with softclipped reads
        if args.boost_ins and "INS" in args.svs_to_assemble:
            logger.info("Generating intervals for insertions")
            assembly_bed = parallel_generate_sc_intervals([args.bam.name], list(contig_whitelist), merged_bed,
                                                          args.workdir,
                                                          num_threads=args.num_threads,
                                                          min_support=args.min_ins_support,
                                                          min_support_frac=args.min_ins_support_frac,
                                                          max_intervals=args.max_ins_intervals, min_mapq=args.min_mapq,
                                                          min_avg_base_qual=args.min_avg_base_qual,
                                                          min_soft_clip=args.min_soft_clip,
                                                          max_soft_clip=args.max_soft_clip,
                                                          max_nm=args.max_nm, min_matches=args.min_matches)
            logger.info("Generated intervals for assembly in %s" % assembly_bed)

        logger.info("Will run assembly now")

        assembled_fasta, ignored_bed = run_spades_parallel(bam=args.bam.name, spades=args.spades, bed=assembly_bed,
                                                           work=spades_tmpdir, pad=SPADES_PAD,
                                                           nthreads=args.num_threads,
                                                           chrs=list(contig_whitelist),
                                                           max_interval_size=args.spades_max_interval_size,
                                                           svs_to_assemble=args.svs_to_assemble,
                                                           stop_on_fail=args.stop_spades_on_fail,
                                                           max_read_pairs=args.extraction_max_read_pairs)
        breakpoints_bed = run_age_parallel(intervals_bed=assembly_bed, reference=args.reference,
                                           assembly=assembled_fasta,
                                           pad=AGE_PAD, age=args.age, chrs=list(contig_whitelist),
                                           nthreads=args.num_threads,
                                           min_contig_len=AGE_MIN_CONTIG_LENGTH, age_workdir=age_tmpdir)

        final_bed = os.path.join(args.workdir, "final.bed")
        if breakpoints_bed:
            if ignored_bed:
                pybedtools.BedTool(breakpoints_bed) \
                    .cat(pybedtools.BedTool(ignored_bed), postmerge=False) \
                    .sort().saveas(final_bed)
            else:
                pybedtools.BedTool(breakpoints_bed).saveas(final_bed)
        elif ignored_bed:
            pybedtools.BedTool(ignored_bed).sort().saveas(final_bed)
        else:
            final_bed = None

        genotyped_bed = parallel_genotype_intervals(final_bed, args.bam.name,
                                                    workdir=os.path.join(args.workdir, "genotyping"),
                                                    nthreads=args.num_threads, chromosomes=list(contig_whitelist),
                                                    window=args.gt_window, isize_mean=args.isize_mean,
                                                    isize_sd=args.isize_sd,
                                                    normal_frac_threshold=args.gt_normal_frac)

        logger.info("Output final VCF file")

        convert_metasv_bed_to_vcf(bedfile=genotyped_bed, vcf_out=final_vcf, sample=args.sample, pass_calls=False)
    else:
        shutil.copy(preasm_vcf, final_vcf)
        pysam.tabix_index(final_vcf, force=True, preset="vcf")

    logger.info("Clean up pybedtools")

    pybedtools.cleanup(remove_all=True)

    logger.info("All Done!")

    return os.EX_OK
<commit_msg>also log arguments<commit_after>from collections import defaultdict
import shutil

from defaults import *
from vcf_utils import *
from sv_interval import SVInterval, get_gaps_file, interval_overlaps_interval_list, merge_intervals
from pindel_reader import PindelReader
from breakdancer_reader import BreakDancerReader
from breakseq_reader import BreakSeqReader
from cnvnator_reader import CNVnatorReader
from generate_sv_intervals import parallel_generate_sc_intervals
from run_spades import run_spades_parallel
from run_age import run_age_parallel
from generate_final_vcf import convert_metasv_bed_to_vcf
from fasta_utils import get_contigs
from genotype import parallel_genotype_intervals

FORMAT = '%(levelname)s %(asctime)-15s %(name)-20s %(message)s'
logging.basicConfig(level=logging.INFO, format=FORMAT)
logger = logging.getLogger(__name__)

def create_dirs(dirlist):
    for dirname in dirlist:
        if not os.path.isdir(dirname):
            logger.info("Creating directory %s" % (dirname))
            os.makedirs(dirname)

def run_metasv(args):
    logger.info("Arguments are " + str(args))

    # Check if there is work to do
    if not (
                                args.pindel_vcf + args.breakdancer_vcf + args.breakseq_vcf + args.cnvnator_vcf + args.pindel_native + args.breakdancer_native + args.breakseq_native + args.cnvnator_native):
        logger.warning("Nothing to merge since no SV file specified")

    # Simple check for arguments
    if not args.disable_assembly:
        if not args.spades:
            logger.error("Spades executable not specified")
            return os.EX_USAGE

        if not args.age:
            logger.error("AGE executable not specified")
            return os.EX_USAGE

    # Create the directories for working
    bedtools_tmpdir = os.path.join(args.workdir, "bedtools")
    create_dirs([args.workdir, args.outdir, bedtools_tmpdir])

    # Reference handling
    if not os.path.isfile(args.reference + ".fai"):
        logger.error("Reference file %s is not indexed" % (args.reference))
        return 1

    fasta_handle = pysam.Fastafile(args.reference) if os.path.isfile(args.reference) else None
    contigs = get_contigs(args.reference)
    include_intervals = sorted(
        [SVInterval(contig.name, 0, contig.length, contig.name, "include", length=contig.length) for contig in contigs])

    # Generate the list of contigs to process
    contig_whitelist = set(args.chromosomes) if args.chromosomes else set([contig.name for contig in contigs])
    if args.keep_standard_contigs:
        contig_whitelist &= set(
            [str(i) for i in xrange(1, 23)] + ["chr%d" % (i) for i in xrange(1, 23)] + ["X", "Y", "MT", "chrX", "chrY",
                                                                                "chrM"])
    logger.info("Only SVs on the following contigs will be reported: %s" % (sorted(list(contig_whitelist))))

    # Load the intervals from different files
    vcf_name_list = [("CNVnator", args.cnvnator_vcf), ("Pindel", args.pindel_vcf),
                     ("BreakDancer", args.breakdancer_vcf),
                     ("BreakSeq", args.breakseq_vcf), ("HaplotypeCaller", args.gatk_vcf)]
    native_name_list = [("CNVnator", args.cnvnator_native, CNVnatorReader),
                        ("Pindel", args.pindel_native, PindelReader),
                        ("BreakSeq", args.breakseq_native, BreakSeqReader),
                        ("BreakDancer", args.breakdancer_native, BreakDancerReader)]

    tools = []
    intervals = {}
    sv_types = set()

    gap_intervals = []
    if args.filter_gaps:
        gaps = args.gaps if args.gaps else get_gaps_file(contig_whitelist)
        gap_intervals = sorted(load_gap_intervals(gaps))

    # Handles native input
    logger.info("Load native files")
    for toolname, nativename, svReader in native_name_list:
        # If no native file is given, ignore the tool
        if not nativename: continue

        tools.append(toolname)
        intervals[toolname] = defaultdict(list)

        for native_file in nativename:
            for record in svReader(native_file, svs_to_report=args.svs_to_report):
                interval = record.to_sv_interval()

                if not interval:
                    # This is the case for SVs we want to skip
                    continue
                if not interval_overlaps_interval_list(interval, gap_intervals) and interval.chrom in contig_whitelist:

                    # Check length
                    if interval.length < args.minsvlen:
                        continue

                    # Set wiggle
                    interval.wiggle = max(args.inswiggle if interval.sv_type == "INS" else 0, args.wiggle)

                    intervals[toolname][interval.sv_type].append(interval)

        sv_types |= set(intervals[toolname].keys())

    # Handles the VCF input cases, we will just deal with these cases
    logger.info("Load VCF files")
    for toolname, vcfname in vcf_name_list:
        # If no VCF is given, ignore the tool
        if not vcfname:
            continue

        tools.append(toolname)
        intervals[toolname] = {}

        vcf_list = []
        for vcffile in vcfname:
            if os.path.isdir(vcffile):
                logger.info("Will load from per-chromosome VCFs from directory %s for tool %s" % (vcffile, toolname))
                vcf_list += [os.path.join(vcffile, "%s.vcf.gz" % contig.name) for contig in contigs if
                             (not contig_whitelist or contig.name in contig_whitelist)]
            else:
                vcf_list.append(vcffile)

        for vcffile in vcf_list:
            load_intervals(vcffile, intervals[toolname], gap_intervals, include_intervals, toolname, contig_whitelist,
                           minsvlen=args.minsvlen, wiggle=args.wiggle, inswiggle=args.inswiggle,
                           svs_to_report=args.svs_to_report)
        sv_types |= set(intervals[toolname].keys())

    logger.info("SV types are %s" % (str(sv_types)))
    tool_merged_intervals = {}
    final_intervals = []

    # This will just output per-tool VCFs, no intra-tool merging is done yet
    if args.enable_per_tool_output:
        logger.info("Output per-tool VCFs")
        for toolname in intervals:
            tool_out = os.path.join(args.outdir, "%s.vcf" % (toolname.lower()))

            logger.info("Outputting single tool VCF for %s" % (str(toolname)))
            vcf_template_reader = vcf.Reader(open(os.path.join(mydir, "resources/template.vcf"), "r"))
            vcf_template_reader.samples = [args.sample]

            intervals_tool = []
            tool_out_fd = open(tool_out, "w")
            vcf_writer = vcf.Writer(tool_out_fd, vcf_template_reader)
            chr_intervals_tool = {contig.name: [] for contig in contigs}
            for sv_type in sv_types:
                if sv_type in intervals[toolname]:
                    intervals_tool.extend([copy.deepcopy(interval) for interval in intervals[toolname][sv_type]])
            for interval in intervals_tool:
                # Marghoob says that this is just to fill-in some metadata
                interval.do_validation(args.overlap_ratio)

                interval.fix_pos()
                chr_intervals_tool[interval.chrom].append(interval)

            for contig in contigs:
                chr_intervals_tool[contig.name].sort()
                for interval in chr_intervals_tool[contig.name]:
                    vcf_record = interval.to_vcf_record(fasta_handle, args.sample)
                    if vcf_record is not None:
                        vcf_writer.write_record(vcf_record)
            tool_out_fd.close()
            vcf_writer.close()
            logger.info("Indexing single tool VCF for %s" % (str(toolname)))
            pysam.tabix_index(tool_out, force=True, preset="vcf")

    # Do merging here
    logger.info("Do merging")
    for sv_type in sv_types:
        logger.info("Processing SVs of type %s" % sv_type)
        tool_merged_intervals[sv_type] = []

        # Do the intra-tool merging
        logger.info("Intra-tool Merging SVs of type %s" % sv_type)
        for tool in tools:
            logger.debug("Is %s in tool keys? %s" % (sv_type, str(intervals[tool].keys())))
            if sv_type not in intervals[tool]:
                logger.debug("%s not in tool %s" % (sv_type, tool))
                continue
            logger.info("First level merging for %s for tool %s" % (sv_type, tool))
            tool_merged_intervals[sv_type] += merge_intervals(intervals[tool][sv_type])

        # Do the inter-tool merging
        logger.info("Inter-tool Merging SVs of type %s" % sv_type)
        merged_intervals = merge_intervals(tool_merged_intervals[sv_type])

        # Intervals which overlap well with merged_intervals
        intervals1 = []
        # Intervals which do not overlap well with merged_intervals.
        # Used to filter out small intervals which got merged with large intervals
        intervals2 = []

        logger.info("Checking overlaps SVs of type %s" % sv_type)
        for interval in tool_merged_intervals[sv_type]:
            if interval_overlaps_interval_list(interval, merged_intervals, args.overlap_ratio, args.overlap_ratio):
                intervals2.append(interval)
            else:
                intervals1.append(interval)
        final_intervals.extend(merge_intervals(intervals1) + merge_intervals(intervals2))

    final_chr_intervals = {contig.name: [] for contig in contigs}
    for interval in final_intervals:
        interval.do_validation(args.overlap_ratio)
        interval.fix_pos()
        if args.minsvlen <= interval.length <= args.maxsvlen:
            final_chr_intervals[interval.chrom].append(interval)

    # This is the merged VCF without assembly, ok for deletions at this point
    logger.info("Output merged VCF without assembly ")
    vcf_template_reader = vcf.Reader(open(os.path.join(mydir, "resources/template.vcf"), "r"))
    vcf_template_reader.samples = [args.sample]
    preasm_vcf = os.path.join(args.workdir, "pre_asm.vcf")
    vcf_fd = open(preasm_vcf, "w")
    vcf_writer = vcf.Writer(vcf_fd, vcf_template_reader)

    final_stats = {}

    bed_intervals = []
    for contig in contigs:
        final_chr_intervals[contig.name].sort()
        for interval in final_chr_intervals[contig.name]:
            vcf_record = interval.to_vcf_record(fasta_handle)
            if vcf_record is not None:
                key = (interval.sv_type, "PASS" if interval.is_validated else "LowQual",
                       "PRECISE" if interval.is_precise else "IMPRECISE", tuple(sorted(list(interval.sources))))
                if key not in final_stats:
                    final_stats[key] = 0
                final_stats[key] += 1
                vcf_writer.write_record(vcf_record)
            bed_interval = interval.to_bed_interval(args.sample)
            if bed_interval is not None:
                bed_intervals.append(bed_interval)
    vcf_fd.close()
    vcf_writer.close()

    # Also save a BED file representation of the merged variants without assembly
    merged_bed = None
    if bed_intervals:
        merged_bed = os.path.join(args.workdir, "metasv.bed")
        pybedtools.BedTool(bed_intervals).saveas(merged_bed)

    for key in sorted(final_stats.keys()):
        logger.info(str(key) + ":" + str(final_stats[key]))

    final_vcf = os.path.join(args.outdir, "variants.vcf")

    # Run assembly here
    if not args.disable_assembly:
        logger.info("Running assembly")

        spades_tmpdir = os.path.join(args.workdir, "spades")
        age_tmpdir = os.path.join(args.workdir, "age")

        create_dirs([spades_tmpdir, age_tmpdir])

        assembly_bed = merged_bed

        # this does the improved assembly location finder with softclipped reads
        if args.boost_ins and "INS" in args.svs_to_assemble:
            logger.info("Generating intervals for insertions")
            assembly_bed = parallel_generate_sc_intervals([args.bam.name], list(contig_whitelist), merged_bed,
                                                          args.workdir,
                                                          num_threads=args.num_threads,
                                                          min_support=args.min_ins_support,
                                                          min_support_frac=args.min_ins_support_frac,
                                                          max_intervals=args.max_ins_intervals, min_mapq=args.min_mapq,
                                                          min_avg_base_qual=args.min_avg_base_qual,
                                                          min_soft_clip=args.min_soft_clip,
                                                          max_soft_clip=args.max_soft_clip,
                                                          max_nm=args.max_nm, min_matches=args.min_matches)
            logger.info("Generated intervals for assembly in %s" % assembly_bed)

        logger.info("Will run assembly now")

        assembled_fasta, ignored_bed = run_spades_parallel(bam=args.bam.name, spades=args.spades, bed=assembly_bed,
                                                           work=spades_tmpdir, pad=SPADES_PAD,
                                                           nthreads=args.num_threads,
                                                           chrs=list(contig_whitelist),
                                                           max_interval_size=args.spades_max_interval_size,
                                                           svs_to_assemble=args.svs_to_assemble,
                                                           stop_on_fail=args.stop_spades_on_fail,
                                                           max_read_pairs=args.extraction_max_read_pairs)
        breakpoints_bed = run_age_parallel(intervals_bed=assembly_bed, reference=args.reference,
                                           assembly=assembled_fasta,
                                           pad=AGE_PAD, age=args.age, chrs=list(contig_whitelist),
                                           nthreads=args.num_threads,
                                           min_contig_len=AGE_MIN_CONTIG_LENGTH, age_workdir=age_tmpdir)

        final_bed = os.path.join(args.workdir, "final.bed")
        if breakpoints_bed:
            if ignored_bed:
                pybedtools.BedTool(breakpoints_bed) \
                    .cat(pybedtools.BedTool(ignored_bed), postmerge=False) \
                    .sort().saveas(final_bed)
            else:
                pybedtools.BedTool(breakpoints_bed).saveas(final_bed)
        elif ignored_bed:
            pybedtools.BedTool(ignored_bed).sort().saveas(final_bed)
        else:
            final_bed = None

        genotyped_bed = parallel_genotype_intervals(final_bed, args.bam.name,
                                                    workdir=os.path.join(args.workdir, "genotyping"),
                                                    nthreads=args.num_threads, chromosomes=list(contig_whitelist),
                                                    window=args.gt_window, isize_mean=args.isize_mean,
                                                    isize_sd=args.isize_sd,
                                                    normal_frac_threshold=args.gt_normal_frac)

        logger.info("Output final VCF file")

        convert_metasv_bed_to_vcf(bedfile=genotyped_bed, vcf_out=final_vcf, sample=args.sample, pass_calls=False)
    else:
        shutil.copy(preasm_vcf, final_vcf)
        pysam.tabix_index(final_vcf, force=True, preset="vcf")

    logger.info("Clean up pybedtools")

    pybedtools.cleanup(remove_all=True)

    logger.info("All Done!")

    return os.EX_OK
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

#add the object name to the download PDF file name, by johnw, 2013/12/29

import simplejson
import time
import base64
import zlib
from openerp.addons.web.http import httprequest
from openerp.addons.web.controllers.main import content_disposition
from openerp.addons.web.controllers.main import Reports as WebReports
content_disposition
@httprequest
def metro_rpt_index(self, req, action, token):
    action = simplejson.loads(action)

    report_srv = req.session.proxy("report")
    context = dict(req.context)
    context.update(action["context"])

    report_data = {}
    report_ids = context["active_ids"]
    if 'report_type' in action:
        report_data['report_type'] = action['report_type']
    if 'datas' in action:
        if 'ids' in action['datas']:
            report_ids = action['datas'].pop('ids')
        report_data.update(action['datas'])

    report_id = report_srv.report(
        req.session._db, req.session._uid, req.session._password,
        action["report_name"], report_ids,
        report_data, context)

    report_struct = None
    while True:
        report_struct = report_srv.report_get(
            req.session._db, req.session._uid, req.session._password, report_id)
        if report_struct["state"]:
            break

        time.sleep(self.POLLING_DELAY)

    report = base64.b64decode(report_struct['result'])
    if report_struct.get('code') == 'zlib':
        report = zlib.decompress(report)
    report_mimetype = self.TYPES_MAPPING.get(
        report_struct['format'], 'octet-stream')
    file_name = action.get('name', 'report')
    ''' comment by john 02/15/2014, to use new logic to get the file name
    if 'name' not in action:
        reports = req.session.model('ir.actions.report.xml')
        res_id = reports.search([('report_name', '=', action['report_name']),],
                                0, False, False, context)
        if len(res_id) > 0:
            #changed by johnw @02/15/2014, use the 'download_filename' field
            file_name = reports.read(res_id[0], ['name'], context)['name']
        else:
            file_name = action['report_name']
    '''
    #new logic to get file name from ir_act_report_xml.download_filename
    #begin
    reports = req.session.model('ir.actions.report.xml')
    res_id = reports.search([('report_name', '=', action['report_name']),],
                            0, False, False, context)
    if len(res_id) > 0:
        names = reports.read(res_id[0], ['name','download_filename'], context)
        if names['download_filename']:
            file_name = names['download_filename']
        else:
            file_name = names['name']
    else:
        file_name = action['report_name']
    #end

    #add the object name to the file name, by johnw, 2013/12/29
    try:
        #only print the name when print one record
        if context['active_ids'] and len(context['active_ids']) == 1:
            model_obj = req.session.model(context['active_model'])
            model_rec_name = model_obj.name_get(context['active_ids'],context)[0][1]
            if model_rec_name:
                file_name = '%s_%s' % (file_name, model_rec_name)
    except Exception:
        pass

    file_name = '%s.%s' % (file_name, report_struct['format'])

    return req.make_response(report,
         headers=[
             ('Content-Disposition', content_disposition(file_name, req)),
             ('Content-Type', report_mimetype),
             ('Content-Length', len(report))],
         cookies={'fileToken': token})

WebReports.index = metro_rpt_index

from osv import fields,osv,orm
class ir_action_report_xml(osv.osv):
    _name="ir.actions.report.xml"
    _inherit ="ir.actions.report.xml"
    _columns={
        'download_filename' : fields.char('Download File Name',size=64),
    }<commit_msg>Fix the issue: ir_sequence.number_next can not be update for the 'standard' sequence<commit_after># -*- coding: utf-8 -*-

#add the object name to the download PDF file name, by johnw, 2013/12/29

import simplejson
import time
import base64
import zlib
from openerp.addons.web.http import httprequest
from openerp.addons.web.controllers.main import content_disposition
from openerp.addons.web.controllers.main import Reports as WebReports
content_disposition
@httprequest
def metro_rpt_index(self, req, action, token):
    action = simplejson.loads(action)

    report_srv = req.session.proxy("report")
    context = dict(req.context)
    context.update(action["context"])

    report_data = {}
    report_ids = context["active_ids"]
    if 'report_type' in action:
        report_data['report_type'] = action['report_type']
    if 'datas' in action:
        if 'ids' in action['datas']:
            report_ids = action['datas'].pop('ids')
        report_data.update(action['datas'])

    report_id = report_srv.report(
        req.session._db, req.session._uid, req.session._password,
        action["report_name"], report_ids,
        report_data, context)

    report_struct = None
    while True:
        report_struct = report_srv.report_get(
            req.session._db, req.session._uid, req.session._password, report_id)
        if report_struct["state"]:
            break

        time.sleep(self.POLLING_DELAY)

    report = base64.b64decode(report_struct['result'])
    if report_struct.get('code') == 'zlib':
        report = zlib.decompress(report)
    report_mimetype = self.TYPES_MAPPING.get(
        report_struct['format'], 'octet-stream')
    file_name = action.get('name', 'report')
    ''' comment by john 02/15/2014, to use new logic to get the file name
    if 'name' not in action:
        reports = req.session.model('ir.actions.report.xml')
        res_id = reports.search([('report_name', '=', action['report_name']),],
                                0, False, False, context)
        if len(res_id) > 0:
            #changed by johnw @02/15/2014, use the 'download_filename' field
            file_name = reports.read(res_id[0], ['name'], context)['name']
        else:
            file_name = action['report_name']
    '''
    #new logic to get file name from ir_act_report_xml.download_filename
    #begin
    reports = req.session.model('ir.actions.report.xml')
    res_id = reports.search([('report_name', '=', action['report_name']),],
                            0, False, False, context)
    if len(res_id) > 0:
        names = reports.read(res_id[0], ['name','download_filename'], context)
        if names['download_filename']:
            file_name = names['download_filename']
        else:
            file_name = names['name']
    else:
        file_name = action['report_name']
    #end

    #add the object name to the file name, by johnw, 2013/12/29
    try:
        #only print the name when print one record
        if context['active_ids'] and len(context['active_ids']) == 1:
            model_obj = req.session.model(context['active_model'])
            model_rec_name = model_obj.name_get(context['active_ids'],context)[0][1]
            if model_rec_name:
                file_name = '%s_%s' % (file_name, model_rec_name)
    except Exception:
        pass

    file_name = '%s.%s' % (file_name, report_struct['format'])

    return req.make_response(report,
         headers=[
             ('Content-Disposition', content_disposition(file_name, req)),
             ('Content-Type', report_mimetype),
             ('Content-Length', len(report))],
         cookies={'fileToken': token})

WebReports.index = metro_rpt_index

'''
fix the next_number updating issue for the standard sequence
'''
from openerp.addons.base.ir.ir_sequence import ir_sequence
def _sequence_next(self, cr, uid, seq_ids, context=None):
    if not seq_ids:
        return False
    if context is None:
        context = {}
    force_company = context.get('force_company')
    if not force_company:
        force_company = self.pool.get('res.users').browse(cr, uid, uid).company_id.id
    sequences = self.read(cr, uid, seq_ids, ['company_id','implementation','number_next','prefix','suffix','padding'])
    preferred_sequences = [s for s in sequences if s['company_id'] and s['company_id'][0] == force_company ]
    seq = preferred_sequences[0] if preferred_sequences else sequences[0]
    if seq['implementation'] == 'standard':
        cr.execute("SELECT nextval('ir_sequence_%03d')" % seq['id'])
        seq['number_next'] = cr.fetchone()
    else:
        cr.execute("SELECT number_next FROM ir_sequence WHERE id=%s FOR UPDATE NOWAIT", (seq['id'],))
#        cr.execute("UPDATE ir_sequence SET number_next=number_next+number_increment WHERE id=%s ", (seq['id'],))
    #fixed by johnw, to update the number_next for all sequence type.
    cr.execute("UPDATE ir_sequence SET number_next=number_next+number_increment WHERE id=%s ", (seq['id'],))
    d = self._interpolation_dict()
    interpolated_prefix = self._interpolate(seq['prefix'], d)
    interpolated_suffix = self._interpolate(seq['suffix'], d)
    return interpolated_prefix + '%%0%sd' % seq['padding'] % seq['number_next'] + interpolated_suffix

ir_sequence._next =  _sequence_next

from osv import fields,osv,orm
class ir_action_report_xml(osv.osv):
    _name="ir.actions.report.xml"
    _inherit ="ir.actions.report.xml"
    _columns={
        'download_filename' : fields.char('Download File Name',size=64),
    }<|endoftext|>


<commit_before>import feedparser
import lxml.html
import xml.etree.ElementTree as ET
from bottle import route, run, view

@route('/api/<url:re:.+>')
def items(url = ''):
	items = {'items' : [], 'url' : url}

	urls = []
	if url: urls.append(url)
	else:
		for channel in channels()['channels']:
			urls.append(channel['url'])

	for url in urls:
		feed = feedparser.parse(url)
		for item in feed.entries:
			items['items'].append({'title': item.title,
							       'description' : item.description,
							       'link' : item.link})
	return items

@route('/api')
def channels():
	tree = ET.parse('subscriptions.xml')
	channels = {'channels' : []}

	for channel in tree.getroot().findall('./body/outline'):
		channels['channels'].append({'title' : channel.get('title'), 'url' : channel.get('xmlUrl')})

	return channels

@route("/")
@route("/<url:re:.+>")
@view('index')
def index(url = ''):
	index = dict(items(url),**channels())
	return index

run(host='localhost', port=3000, reloader = True, debug = True)
<commit_msg>Peewee orm.<commit_after>import feedparser
import lxml.html
import xml.etree.ElementTree as ET
from bottle import route, run, view
import peewee
from peewee import *

class Channel(peewee.Model):
	title = CharField()
	url = CharField()

class Item(peewee.Model):
	title = CharField()
	description = CharField()
	url = CharField()
	channel = peewee.ForeignKeyField(Channel)

@route('/api/<url:re:.+>')
def items(url = ''):
	items = {'items' : [], 'url' : url}

	urls = []
	if url: urls.append(url)
	else:
		for channel in channels()['channels']:
			urls.append(channel['url'])

	for url in urls:
		feed = feedparser.parse(url)
		for item in feed.entries:
			items['items'].append({'title': item.title,
							       'description' : item.description,
							       'link' : item.link})
	return items

@route('/api')
def channels():
	tree = ET.parse('subscriptions.xml')
	channels = {'channels' : []}

	for channel in tree.getroot().findall('./body/outline'):
		channels['channels'].append({'title' : channel.get('title'), 'url' : channel.get('xmlUrl')})

	return channels

@route("/")
@route("/<url:re:.+>")
@view('index')
def index(url = ''):
	index = dict(items(url),**channels())
	return index

@route('/bla')
def bla():
	return Channel;

run(host='localhost', port=3000, reloader = True, debug = True)
<|endoftext|>


<commit_before># http://ls.pwd.io/2014/08/singly-and-doubly-linked-lists-in-python/

class Node(object):
    def __init__(self, data):
        self.next = None
        self.prev = None
        self.data = data

class DoublyList(object):
    def __init__(self):
        self.head = None
        self.tail = None

    def insert(self, data):
        """Insert data to at head"""
        node = Node(data)
        if self.head is None:
            self.head = self.tail = node
        elif self.head is not None:
            node.prev = self.head
            self.head.prev = node
            self.head = node

    def append(self, data):
        """Insert data to at head"""
        node = Node(data)
        if self.head is None:
            self.head = self.tail = node
        elif self.head is not None:
            node.next = self.tail
            self.tail.next = node
            self.tail = node

    def pop(self):
        """Remove and return head value"""
        if self.head is not None:
            value = self.head.data
            self.head = self.head.next
            if self.head is not None:
                self.head.prev = None
        return value

    def shift(self):
        """Remove and return tail value"""
        if self.tail is not None:
            value = self.tail.data
            self.tail = self.tail.prev
            if self.tail is not None:
                self.tail.next = None
        return value

    def remove(self, data):
        """Remove first instance of val"""
        curr_node = self.head
        while curr_node is not None:
            if curr_node.data == data:
                #want to remove
                prev_node = curr_node.prev
                next_node = curr_node.next
                #first element in list
                if self.head == curr_node:
                    self.head = curr_node.next
                elif self.tail == curr_node:
                    self.tail = curr_node.prev
                else:
                    prev_node.next = next_node
                    next_node.prev = prev_node
                    curr_node = curr_node.next
        return LookupError("Data not in list.")

    #     def search(self, data):
    #     find_node = self.head
    #     while find_node is not None:
    #         if find_node.data == data:
    #             return find_node
    #         else:
    #             find_node = find_node.next
    #
    # def remove(self, node):
    #     """Remove node from list"""
    #     prev_node = None
    #     cur_node = self.head
    #     while cur_node is not None:
    #         if cur_node == node:
    #             if prev_node is None:
    #                 self.head = node.next
    #                 return cur_node
    #             else:
    #                 prev_node.next = cur_node.next
    #             return cur_node
    #         else:
    #             cur_node = cur_node.next
    #     return None
<commit_msg>fixed remove method<commit_after># http://ls.pwd.io/2014/08/singly-and-doubly-linked-lists-in-python/

class Node(object):
    def __init__(self, data):
        self.next = None
        self.prev = None
        self.data = data

class DoublyList(object):
    def __init__(self):
        self.head = None
        self.tail = None

    def insert(self, data):
        """Insert data to at head"""
        node = Node(data)
        if self.head is None:
            self.head = self.tail = node
        elif self.head is not None:
            node.prev = self.head
            self.head.prev = node
            self.head = node

    def append(self, data):
        """Insert data to at head"""
        node = Node(data)
        if self.head is None:
            self.head = self.tail = node
        elif self.head is not None:
            node.next = self.tail
            self.tail.next = node
            self.tail = node

    def pop(self):
        """Remove and return head value"""
        if self.head is not None:
            value = self.head.data
            self.head = self.head.next
            if self.head is not None:
                self.head.prev = None
        return value

    def shift(self):
        """Remove and return tail value"""
        if self.tail is not None:
            value = self.tail.data
            self.tail = self.tail.prev
            if self.tail is not None:
                self.tail.next = None
        return value

    def remove(self, data):
        """Remove first instance of val"""
        if self.head is None:
            raise LookupError('Data not in list.')
        elif self.head.data == data:
            self.head = self.head.next
        elif self.tail.data == data:
            self.tail = self.tail.prev
        else:
            curr_node = self.head
            while True:
                if curr_node.data == data:
                    curr_node.prev.next = curr_node.next
                    curr_node.next.prev = curr_node.prev
                    break
                else:
                    if curr_node.next is None:
                        raise LookupError('Data not in list.')
                    else:
                        curr_node = curr_node.next
<|endoftext|>


<commit_before>import logging
import re
import networkx as nx
from dowhy.utils.api import parse_state
import itertools

class CausalGraph:

    """Class for creating and modifying the causal graph.

       Accepts a graph string (or a text file) in gml format (preferred) and dot format. Graphviz-like attributes can be set for edges and nodes. E.g. style="dashed" as an edge attribute ensures that the edge is drawn with a dashed line.

        If a graph string is not given, names of treatment, outcome, and confounders, instruments and effect modifiers (if any) can be provided to create the graph.
    """

    def __init__(self,
                 treatment_name, outcome_name,
                 graph=None,
                 common_cause_names=None,
                 instrument_names=None,
                 effect_modifier_names=None,
                 observed_node_names=None,
                 missing_nodes_as_confounders=False):
        self.treatment_name = parse_state(treatment_name)
        self.outcome_name = parse_state(outcome_name)
        instrument_names = parse_state(instrument_names)
        common_cause_names = parse_state(common_cause_names)
        self.logger = logging.getLogger(__name__)

        if graph is None:
            self._graph = nx.DiGraph()
            self._graph = self.build_graph(common_cause_names,
                                           instrument_names, effect_modifier_names)
        elif re.match(r".*\.dot", graph):
            # load dot file
            try:
                import pygraphviz as pgv
                self._graph = nx.DiGraph(nx.drawing.nx_agraph.read_dot(graph))
            except Exception as e:
                self.logger.error("Pygraphviz cannot be loaded. " + str(e) + "\nTrying pydot...")
                try:
                    import pydot
                    self._graph = nx.DiGraph(nx.drawing.nx_pydot.read_dot(graph))
                except Exception as e:
                    self.logger.error("Error: Pydot cannot be loaded. " + str(e))
                    raise e
        elif re.match(r".*\.gml", graph):
            self._graph = nx.DiGraph(nx.read_gml(graph))
        elif re.match(r".*graph\s*\{.*\}\s*", graph):
            try:
                import pygraphviz as pgv
                self._graph = pgv.AGraph(graph, strict=True, directed=True)
                self._graph = nx.drawing.nx_agraph.from_agraph(self._graph)
            except Exception as e:
                self.logger.error("Error: Pygraphviz cannot be loaded. " + str(e) + "\nTrying pydot ...")
                try:
                    import pydot
                    P_list = pydot.graph_from_dot_data(graph)
                    self._graph = nx.drawing.nx_pydot.from_pydot(P_list[0])
                except Exception as e:
                    self.logger.error("Error: Pydot cannot be loaded. " + str(e))
                    raise e
        elif re.match(".*graph\s*\[.*\]\s*", graph):
            self._graph = nx.DiGraph(nx.parse_gml(graph))
        else:
            self.logger.error("Error: Please provide graph (as string or text file) in dot or gml format.")
            self.logger.error("Error: Incorrect graph format")
            raise ValueError
        if missing_nodes_as_confounders:
            self._graph = self.add_missing_nodes_as_common_causes(observed_node_names)
        # Adding node attributes
        self._graph = self.add_node_attributes(observed_node_names)
        #TODO do not add it here. CausalIdentifier should call causal_graph to add an unobserved common cause if needed. This also ensures that we do not need get_common_causes in this class.
        self._graph = self.add_unobserved_common_cause(observed_node_names)

    def view_graph(self, layout="dot"):
        out_filename = "causal_model.png"
        try:
            import pygraphviz as pgv
            agraph = nx.drawing.nx_agraph.to_agraph(self._graph)
            agraph.draw(out_filename, format="png", prog=layout)
        except:
            self.logger.warning("Warning: Pygraphviz cannot be loaded. Check that graphviz and pygraphviz are installed.")
            self.logger.info("Using Matplotlib for plotting")
            import matplotlib.pyplot as plt
            solid_edges = [(n1,n2) for n1,n2, e in self._graph.edges(data=True) if 'style' not in e ]
            dashed_edges =[(n1,n2) for n1,n2, e in self._graph.edges(data=True) if ('style' in e and e['style']=="dashed") ]
            plt.clf()
            pos = nx.layout.shell_layout(self._graph)
            nodes = nx.draw_networkx_nodes(self._graph, pos)
            edges = nx.draw_networkx_edges(
                    self._graph,
                    pos,
                    edgelist=solid_edges,
                    arrowstyle="fancy",
                    arrowsize=1)
            edges = nx.draw_networkx_edges(
                    self._graph,
                    pos,
                    edgelist=dashed_edges,
                    arrowstyle="->",
                    style="dashed",
                    arrowsize=2)
            labels = nx.draw_networkx_labels(self._graph, pos)
            #nx.draw_networkx(self._graph, pos=nx.shell_layout(self._graph))
            plt.axis('off')
            plt.savefig(out_filename)
            plt.draw()

    def build_graph(self, common_cause_names, instrument_names, effect_modifier_names):
        """ Creates nodes and edges based on variable names and their semantics.

        Currently only considers the graphical representation of "direct" effect modifiers. Thus, all effect modifiers are assumed to be "direct" unless otherwise expressed using a graph. Based on the taxonomy of effect modifiers by VanderWheele and Robins: "Four types of effect modification: A classification based on directed acyclic graphs. Epidemiology. 2007."
        """

        for treatment in self.treatment_name:
            self._graph.add_node(treatment, observed="yes")
        for outcome in self.outcome_name:
            self._graph.add_node(outcome, observed="yes")
        for treatment, outcome in itertools.product(self.treatment_name, self.outcome_name):
            self._graph.add_edge(treatment, outcome)

        # Adding common causes
        if common_cause_names is not None:
            for node_name in common_cause_names:
                for treatment, outcome in itertools.product(self.treatment_name, self.outcome_name):
                    self._graph.add_node(node_name, observed="yes")
                    self._graph.add_edge(node_name, treatment)
                    self._graph.add_edge(node_name, outcome)

        # Adding instruments
        if instrument_names:
            if type(instrument_names[0]) != tuple:
                if len(self.treatment_name) > 1:
                    self.logger.info("Assuming Instrument points to all treatments! Use tuples for more granularity.")
                for instrument, treatment in itertools.product(instrument_names, self.treatment_name):
                    self._graph.add_node(instrument, observed="yes")
                    self._graph.add_edge(instrument, treatment)
            else:
                for instrument, treatment in itertools.product(instrument_names):
                    self._graph.add_node(instrument, observed="yes")
                    self._graph.add_edge(instrument, treatment)

        # Adding effect modifiers
        if effect_modifier_names is not None:
            for node_name in effect_modifier_names:
                if node_name not in common_cause_names:
                    for outcome in self.outcome_name:
                        self._graph.add_node(node_name, observed="yes")
                        self._graph.add_edge(node_name, outcome, style = "dotted", headport="s", tailport="n")
                        self._graph.add_edge(outcome, node_name, style = "dotted", headport="n", tailport="s") # TODO make the ports more general so that they apply not just to top-bottom node configurations
        return self._graph

    def add_node_attributes(self, observed_node_names):
        for node_name in self._graph:
            if node_name in observed_node_names:
                self._graph.nodes[node_name]["observed"] = "yes"
            else:
                self._graph.nodes[node_name]["observed"] = "no"
        return self._graph

    def add_missing_nodes_as_common_causes(self, observed_node_names):
        # Adding columns in the dataframe as confounders that were not in the graph
        for node_name in observed_node_names:
            if node_name not in self._graph:
                self._graph.add_node(node_name, observed="yes")
                for treatment_outcome_node in self.treatment_name + self.outcome_name:
                    self._graph.add_edge(node_name, treatment_outcome_node)
        return self._graph

    def add_unobserved_common_cause(self, observed_node_names):
        # Adding unobserved confounders
        current_common_causes = self.get_common_causes(self.treatment_name,
                                                       self.outcome_name)
        create_new_common_cause = True
        for node_name in current_common_causes:
            if self._graph.nodes[node_name]["observed"] == "no":
                create_new_common_cause = False
        if create_new_common_cause:
            uc_label = "Unobserved Confounders"
            self._graph.add_node('U', label=uc_label, observed="no")
            for node in self.treatment_name + self.outcome_name:
                self._graph.add_edge('U', node)
            self.logger.info('If this is observed data (not from a randomized experiment), there might always be missing confounders. Adding a node named "Unobserved Confounders" to reflect this.')
        return self._graph

    def get_unconfounded_observed_subgraph(self):
        observed_nodes = [node for node in self._graph.nodes() if self._graph.nodes[node]["observed"] == "yes"]
        return self._graph.subgraph(observed_nodes)

    def do_surgery(self, node_names, remove_outgoing_edges=False,
                   remove_incoming_edges=False):
        new_graph = self._graph.copy()
        for node_name in node_names:
            if remove_outgoing_edges:
                children = new_graph.successors(node_name)
                edges_bunch = [(node_name, child) for child in children]
                new_graph.remove_edges_from(edges_bunch)
            if remove_incoming_edges:
                parents = new_graph.predecessors(node_name)
                edges_bunch = [(parent, node_name) for parent in parents]
                new_graph.remove_edges_from(edges_bunch)
        return new_graph

    def get_causes(self, nodes, remove_edges = None):
        nodes = parse_state(nodes)
        new_graph=None
        if remove_edges is not None:
            new_graph = self._graph.copy()  # caution: shallow copy of the attributes
            sources = parse_state(remove_edges["sources"])
            targets = parse_state(remove_edges["targets"])
            for s in sources:
                for t in targets:
                    new_graph.remove_edge(s, t)
        causes = set()
        for v in nodes:
            causes = causes.union(self.get_ancestors(v, new_graph=new_graph))
        return causes

    def get_common_causes(self, nodes1, nodes2):
        """
        Assume that nodes1 causes nodes2 (e.g., nodes1 are the treatments and nodes2 are the outcomes)
        """
        # TODO Refactor to remove this from here and only implement this logic in causalIdentifier. Unnecessary assumption of nodes1 to be causing nodes2.
        nodes1 = parse_state(nodes1)
        nodes2 = parse_state(nodes2)
        causes_1 = set()
        causes_2 = set()
        for node in nodes1:
            causes_1 = causes_1.union(self.get_ancestors(node))
        for node in nodes2:
            # Cannot simply compute ancestors, since that will also include nodes1 and its parents (e.g. instruments)
            parents_2 = self.get_parents(node)
            for parent in parents_2:
                if parent not in nodes1:
                    causes_2 = causes_2.union(set([parent,]))
                    causes_2 = causes_2.union(self.get_ancestors(parent))
        return list(causes_1.intersection(causes_2))

    def get_effect_modifiers(self, nodes1, nodes2):
        modifiers = set()
        for node in nodes2:
            modifiers = modifiers.union(self.get_ancestors(node))
        modifiers = modifiers.difference(nodes1)
        for node in nodes1:
            modifiers = modifiers.difference(self.get_ancestors(node))
        return list(modifiers)

    def get_parents(self, node_name):
        return set(self._graph.predecessors(node_name))

    def get_ancestors(self, node_name, new_graph=None):
        if new_graph is None:
            graph=self._graph
        else:
            graph=new_graph
        return set(nx.ancestors(graph, node_name))

    def get_descendants(self, node_name):
        return set(nx.descendants(self._graph, node_name))

    def all_observed(self, node_names):
        for node_name in node_names:
            if self._graph.nodes[node_name]["observed"] != "yes":
                return False

        return True

    def filter_unobserved_variables(self, node_names):
        observed_node_names = list()
        for node_name in node_names:
            if self._graph.nodes[node_name]["observed"] == "yes":
                observed_node_names.append(node_name)

        return observed_node_names

    def get_instruments(self, treatment_nodes, outcome_nodes):
        treatment_nodes = parse_state(treatment_nodes)
        outcome_nodes = parse_state(outcome_nodes)
        parents_treatment = set()
        for node in treatment_nodes:
            parents_treatment = parents_treatment.union(self.get_parents(node))
        g_no_parents_treatment = self.do_surgery(treatment_nodes,
                                                 remove_incoming_edges=True)
        ancestors_outcome = set()
        for node in outcome_nodes:
            ancestors_outcome = ancestors_outcome.union(nx.ancestors(g_no_parents_treatment, node))

        # [TODO: double check these work with multivariate implementation:]
        # Exclusion
        candidate_instruments = parents_treatment.difference(ancestors_outcome)
        self.logger.debug("Candidate instruments after exclusion %s",
                          candidate_instruments)
        # As-if-random setup
        children_causes_outcome = [nx.descendants(g_no_parents_treatment, v)
                                   for v in ancestors_outcome]
        children_causes_outcome = set([item
                                       for sublist in children_causes_outcome
                                       for item in sublist])

        # As-if-random
        instruments = candidate_instruments.difference(children_causes_outcome)
        return list(instruments)
<commit_msg>fix networkx graph issues (#142)<commit_after>import logging
import re
import networkx as nx
from dowhy.utils.api import parse_state
import itertools

class CausalGraph:

    """Class for creating and modifying the causal graph.

       Accepts a graph string (or a text file) in gml format (preferred) and dot format. Graphviz-like attributes can be set for edges and nodes. E.g. style="dashed" as an edge attribute ensures that the edge is drawn with a dashed line.

        If a graph string is not given, names of treatment, outcome, and confounders, instruments and effect modifiers (if any) can be provided to create the graph.
    """

    def __init__(self,
                 treatment_name, outcome_name,
                 graph=None,
                 common_cause_names=None,
                 instrument_names=None,
                 effect_modifier_names=None,
                 observed_node_names=None,
                 missing_nodes_as_confounders=False):
        self.treatment_name = parse_state(treatment_name)
        self.outcome_name = parse_state(outcome_name)
        instrument_names = parse_state(instrument_names)
        common_cause_names = parse_state(common_cause_names)
        self.logger = logging.getLogger(__name__)

        if graph is None:
            self._graph = nx.DiGraph()
            self._graph = self.build_graph(common_cause_names,
                                           instrument_names, effect_modifier_names)
        elif re.match(r".*\.dot", graph):
            # load dot file
            try:
                import pygraphviz as pgv
                self._graph = nx.DiGraph(nx.drawing.nx_agraph.read_dot(graph))
            except Exception as e:
                self.logger.error("Pygraphviz cannot be loaded. " + str(e) + "\nTrying pydot...")
                try:
                    import pydot
                    self._graph = nx.DiGraph(nx.drawing.nx_pydot.read_dot(graph))
                except Exception as e:
                    self.logger.error("Error: Pydot cannot be loaded. " + str(e))
                    raise e
        elif re.match(r".*\.gml", graph):
            self._graph = nx.DiGraph(nx.read_gml(graph))
        elif re.match(r".*graph\s*\{.*\}\s*", graph):
            try:
                import pygraphviz as pgv
                self._graph = pgv.AGraph(graph, strict=True, directed=True)
                self._graph = nx.drawing.nx_agraph.from_agraph(self._graph)
            except Exception as e:
                self.logger.error("Error: Pygraphviz cannot be loaded. " + str(e) + "\nTrying pydot ...")
                try:
                    import pydot
                    P_list = pydot.graph_from_dot_data(graph)
                    self._graph = nx.drawing.nx_pydot.from_pydot(P_list[0])
                except Exception as e:
                    self.logger.error("Error: Pydot cannot be loaded. " + str(e))
                    raise e
        elif re.match(".*graph\s*\[.*\]\s*", graph):
            self._graph = nx.DiGraph(nx.parse_gml(graph))
        else:
            self.logger.error("Error: Please provide graph (as string or text file) in dot or gml format.")
            self.logger.error("Error: Incorrect graph format")
            raise ValueError
        if missing_nodes_as_confounders:
            self._graph = self.add_missing_nodes_as_common_causes(observed_node_names)
        # Adding node attributes
        self._graph = self.add_node_attributes(observed_node_names)
        #TODO do not add it here. CausalIdentifier should call causal_graph to add an unobserved common cause if needed. This also ensures that we do not need get_common_causes in this class.
        self._graph = self.add_unobserved_common_cause(observed_node_names)

    def view_graph(self, layout="dot"):
        out_filename = "causal_model.png"
        try:
            import pygraphviz as pgv
            agraph = nx.drawing.nx_agraph.to_agraph(self._graph)
            agraph.draw(out_filename, format="png", prog=layout)
        except:
            self.logger.warning("Warning: Pygraphviz cannot be loaded. Check that graphviz and pygraphviz are installed.")
            self.logger.info("Using Matplotlib for plotting")
            import matplotlib.pyplot as plt

            solid_edges = [(n1,n2) for n1,n2, e in self._graph.edges(data=True) if 'style' not in e ]
            dashed_edges =[(n1,n2) for n1,n2, e in self._graph.edges(data=True) if ('style' in e and e['style']=="dashed") ]
            plt.clf()

            pos = nx.layout.shell_layout(self._graph)
            nx.draw_networkx_nodes(self._graph, pos, node_color='yellow',node_size=400 )
            nx.draw_networkx_edges(
                    self._graph,
                    pos,
                    edgelist=solid_edges,
                    arrowstyle="-|>",
                    arrowsize=12)
            nx.draw_networkx_edges(
                    self._graph,
                    pos,
                    edgelist=dashed_edges,
                    arrowstyle="-|>",
                    style="dashed",
                    arrowsize=12)

            labels = nx.draw_networkx_labels(self._graph, pos)

            plt.axis('off')
            plt.savefig(out_filename)
            plt.draw()

    def build_graph(self, common_cause_names, instrument_names, effect_modifier_names):
        """ Creates nodes and edges based on variable names and their semantics.

        Currently only considers the graphical representation of "direct" effect modifiers. Thus, all effect modifiers are assumed to be "direct" unless otherwise expressed using a graph. Based on the taxonomy of effect modifiers by VanderWheele and Robins: "Four types of effect modification: A classification based on directed acyclic graphs. Epidemiology. 2007."
        """

        for treatment in self.treatment_name:
            self._graph.add_node(treatment, observed="yes")
        for outcome in self.outcome_name:
            self._graph.add_node(outcome, observed="yes")
        for treatment, outcome in itertools.product(self.treatment_name, self.outcome_name):
            self._graph.add_edge(treatment, outcome)

        # Adding common causes
        if common_cause_names is not None:
            for node_name in common_cause_names:
                for treatment, outcome in itertools.product(self.treatment_name, self.outcome_name):
                    self._graph.add_node(node_name, observed="yes")
                    self._graph.add_edge(node_name, treatment)
                    self._graph.add_edge(node_name, outcome)

        # Adding instruments
        if instrument_names:
            if type(instrument_names[0]) != tuple:
                if len(self.treatment_name) > 1:
                    self.logger.info("Assuming Instrument points to all treatments! Use tuples for more granularity.")
                for instrument, treatment in itertools.product(instrument_names, self.treatment_name):
                    self._graph.add_node(instrument, observed="yes")
                    self._graph.add_edge(instrument, treatment)
            else:
                for instrument, treatment in itertools.product(instrument_names):
                    self._graph.add_node(instrument, observed="yes")
                    self._graph.add_edge(instrument, treatment)

        # Adding effect modifiers
        if effect_modifier_names is not None:
            for node_name in effect_modifier_names:
                if node_name not in common_cause_names:
                    for outcome in self.outcome_name:
                        self._graph.add_node(node_name, observed="yes")
                        self._graph.add_edge(node_name, outcome, style = "dotted", headport="s", tailport="n")
                        self._graph.add_edge(outcome, node_name, style = "dotted", headport="n", tailport="s") # TODO make the ports more general so that they apply not just to top-bottom node configurations
        return self._graph

    def add_node_attributes(self, observed_node_names):
        for node_name in self._graph:
            if node_name in observed_node_names:
                self._graph.nodes[node_name]["observed"] = "yes"
            else:
                self._graph.nodes[node_name]["observed"] = "no"
        return self._graph

    def add_missing_nodes_as_common_causes(self, observed_node_names):
        # Adding columns in the dataframe as confounders that were not in the graph
        for node_name in observed_node_names:
            if node_name not in self._graph:
                self._graph.add_node(node_name, observed="yes")
                for treatment_outcome_node in self.treatment_name + self.outcome_name:
                    self._graph.add_edge(node_name, treatment_outcome_node)
        return self._graph

    def add_unobserved_common_cause(self, observed_node_names):
        # Adding unobserved confounders
        current_common_causes = self.get_common_causes(self.treatment_name,
                                                       self.outcome_name)
        create_new_common_cause = True
        for node_name in current_common_causes:
            if self._graph.nodes[node_name]["observed"] == "no":
                create_new_common_cause = False
        if create_new_common_cause:
            uc_label = "Unobserved Confounders"
            self._graph.add_node('U', label=uc_label, observed="no")
            for node in self.treatment_name + self.outcome_name:
                self._graph.add_edge('U', node)
            self.logger.info('If this is observed data (not from a randomized experiment), there might always be missing confounders. Adding a node named "Unobserved Confounders" to reflect this.')
        return self._graph

    def get_unconfounded_observed_subgraph(self):
        observed_nodes = [node for node in self._graph.nodes() if self._graph.nodes[node]["observed"] == "yes"]
        return self._graph.subgraph(observed_nodes)

    def do_surgery(self, node_names, remove_outgoing_edges=False,
                   remove_incoming_edges=False):
        new_graph = self._graph.copy()
        for node_name in node_names:
            if remove_outgoing_edges:
                children = new_graph.successors(node_name)
                edges_bunch = [(node_name, child) for child in children]
                new_graph.remove_edges_from(edges_bunch)
            if remove_incoming_edges:
                parents = new_graph.predecessors(node_name)
                edges_bunch = [(parent, node_name) for parent in parents]
                new_graph.remove_edges_from(edges_bunch)
        return new_graph

    def get_causes(self, nodes, remove_edges = None):
        nodes = parse_state(nodes)
        new_graph=None
        if remove_edges is not None:
            new_graph = self._graph.copy()  # caution: shallow copy of the attributes
            sources = parse_state(remove_edges["sources"])
            targets = parse_state(remove_edges["targets"])
            for s in sources:
                for t in targets:
                    new_graph.remove_edge(s, t)
        causes = set()
        for v in nodes:
            causes = causes.union(self.get_ancestors(v, new_graph=new_graph))
        return causes

    def get_common_causes(self, nodes1, nodes2):
        """
        Assume that nodes1 causes nodes2 (e.g., nodes1 are the treatments and nodes2 are the outcomes)
        """
        # TODO Refactor to remove this from here and only implement this logic in causalIdentifier. Unnecessary assumption of nodes1 to be causing nodes2.
        nodes1 = parse_state(nodes1)
        nodes2 = parse_state(nodes2)
        causes_1 = set()
        causes_2 = set()
        for node in nodes1:
            causes_1 = causes_1.union(self.get_ancestors(node))
        for node in nodes2:
            # Cannot simply compute ancestors, since that will also include nodes1 and its parents (e.g. instruments)
            parents_2 = self.get_parents(node)
            for parent in parents_2:
                if parent not in nodes1:
                    causes_2 = causes_2.union(set([parent,]))
                    causes_2 = causes_2.union(self.get_ancestors(parent))
        return list(causes_1.intersection(causes_2))

    def get_effect_modifiers(self, nodes1, nodes2):
        modifiers = set()
        for node in nodes2:
            modifiers = modifiers.union(self.get_ancestors(node))
        modifiers = modifiers.difference(nodes1)
        for node in nodes1:
            modifiers = modifiers.difference(self.get_ancestors(node))
        return list(modifiers)

    def get_parents(self, node_name):
        return set(self._graph.predecessors(node_name))

    def get_ancestors(self, node_name, new_graph=None):
        if new_graph is None:
            graph=self._graph
        else:
            graph=new_graph
        return set(nx.ancestors(graph, node_name))

    def get_descendants(self, node_name):
        return set(nx.descendants(self._graph, node_name))

    def all_observed(self, node_names):
        for node_name in node_names:
            if self._graph.nodes[node_name]["observed"] != "yes":
                return False

        return True

    def filter_unobserved_variables(self, node_names):
        observed_node_names = list()
        for node_name in node_names:
            if self._graph.nodes[node_name]["observed"] == "yes":
                observed_node_names.append(node_name)

        return observed_node_names

    def get_instruments(self, treatment_nodes, outcome_nodes):
        treatment_nodes = parse_state(treatment_nodes)
        outcome_nodes = parse_state(outcome_nodes)
        parents_treatment = set()
        for node in treatment_nodes:
            parents_treatment = parents_treatment.union(self.get_parents(node))
        g_no_parents_treatment = self.do_surgery(treatment_nodes,
                                                 remove_incoming_edges=True)
        ancestors_outcome = set()
        for node in outcome_nodes:
            ancestors_outcome = ancestors_outcome.union(nx.ancestors(g_no_parents_treatment, node))

        # [TODO: double check these work with multivariate implementation:]
        # Exclusion
        candidate_instruments = parents_treatment.difference(ancestors_outcome)
        self.logger.debug("Candidate instruments after exclusion %s",
                          candidate_instruments)
        # As-if-random setup
        children_causes_outcome = [nx.descendants(g_no_parents_treatment, v)
                                   for v in ancestors_outcome]
        children_causes_outcome = set([item
                                       for sublist in children_causes_outcome
                                       for item in sublist])

        # As-if-random
        instruments = candidate_instruments.difference(children_causes_outcome)
        return list(instruments)
<|endoftext|>


<commit_before>#
# Copyright (c) 2005-2007 rPath, Inc.
#
# All rights reserved
#

import re, os, sys
import pwd
from mint import config
from mint import projects
from mint import schema
from conary.lib import util
from conary import dbstore
from conary import versions
import conary.errors
import conary.server.schema
from conary.conaryclient import cmdline

schemaCutoff = 37
knownGroupVersions = ('3\.1\.5.*', '4\..*')

staticPaths = ['config', 'entitlements', 'logs', 'installable_iso.conf',
               'toolkit', 'iso_gen.conf', 'live_iso.conf',
               'bootable_image.conf']
# toolkit is linked to outgoing job server architecture.

def backup(cfg, out, backupMirrors = False):
    reposContentsDir = cfg.reposContentsDir.split(' ')
    backupPath = os.path.join(cfg.dataPath, 'tmp', 'backup')
    print >> out, backupPath
    util.mkdirChain(backupPath)
    dumpPath = os.path.join(backupPath, 'db.dump')
    if cfg.dbDriver == 'sqlite':
        util.execute("echo '.dump' | sqlite3 %s > %s" % (cfg.dbPath, dumpPath))
    extraArgs = ""
    if not backupMirrors:
        extraArgs = " WHERE NOT external"
    if cfg.reposDBDriver == 'sqlite':
        db = dbstore.connect(cfg.dbPath, cfg.dbDriver)
        cu = db.cursor()
        cu.execute("SELECT hostname, domainname FROM Projects%s" % extraArgs)
        for reposDir in ['.'.join(x) for x in cu.fetchall()]:
            cu.execute("SELECT toName FROM RepNameMap WHERE fromName=?",
                    reposDir)
            res = cu.fetchone()
            reposDir = res and res[0] or reposDir
            reposDBPath = cfg.reposDBPath % reposDir
            dumpPath = os.path.join(backupPath, '%s.dump' % reposDir)
            if os.path.exists(reposDBPath):
                util.execute("echo '.dump' | sqlite3 %s > %s" % \
                              (reposDBPath, dumpPath))
                print >> out, reposContentsDir[0] % reposDir
    elif cfg.reposDBDriver == 'postgresql':
        db = dbstore.connect(cfg.dbPath, cfg.dbDriver)
        cu = db.cursor()
        cu.execute("SELECT hostname, domainname FROM Projects%s" % extraArgs)
        for reposDir in ['.'.join(x) for x in cu.fetchall()]:
            reposDbName = reposDir.translate(projects.transTables['postgresql'])
            dbUser = cfg.reposDBPath.split('@')[0]
            dumpPath = os.path.join(backupPath, '%s.dump' % reposDbName)
            rdb = dbstore.connect(cfg.reposDBPath % 'postgres', 'postgresql')
            rcu = rdb.cursor()
            rcu.execute("SELECT datname FROM pg_database WHERE datname=?",
                         reposDbName)
            if rcu.fetchone():
                util.execute("pg_dump -U %s -c -O -d %s > %s" %\
                             (dbUser, reposDbName, dumpPath))
                print >> out, reposContentsDir[0] % reposDir
            rdb.close()

    for d in staticPaths:
        path = os.path.join(cfg.dataPath, d)
        if os.path.exists(path):
            print >> out, path

def restore(cfg):
    backupPath = os.path.join(cfg.dataPath, 'tmp', 'backup')
    dumpPath = os.path.join(backupPath, 'db.dump')
    if cfg.dbDriver == 'sqlite':
        if os.path.exists(cfg.dbPath):
            os.unlink(cfg.dbPath)
        util.execute("sqlite3 %s < %s" % (cfg.dbPath, dumpPath))
    db = dbstore.connect(cfg.dbPath, cfg.dbDriver)
    schema.loadSchema(db, cfg)
    cu = db.cursor()
    cu.execute("SELECT hostname, domainname, projectId FROM Projects")
    for hostname, domainname, projectId in [x for x in cu.fetchall()]:
        repo = hostname + '.' + domainname
        cu.execute('SELECT toName FROM RepNameMap WHERE fromName=?', repo)
        r = cu.fetchone()
        if r:
            repo = r[0]
        dumpPath =  os.path.join(backupPath, repo).translate(projects.transTables[cfg.reposDBDriver]) + '.dump'
        reposPath = os.path.join(cfg.reposPath, repo)
        if os.path.exists(dumpPath):
            util.mkdirChain(os.path.join(reposPath, 'tmp'))
            if cfg.reposDBDriver == 'sqlite':
                dbPath = cfg.reposDBPath % repo
                if os.path.exists(dbPath):
                    os.unlink(dbPath)
                util.execute('sqlite3 %s < %s' % (dbPath, dumpPath))
                rdb = dbstore.connect(dbPath, cfg.reposDBDriver)
                conary.server.schema.loadSchema(rdb)
                rdb.close()
            elif cfg.reposDBDriver == 'postgresql':
                pgRepo = repo.translate(projects.transTables['postgresql'])
                dbPath = cfg.reposDBPath % pgRepo
                rdb = dbstore.connect('%s' % (cfg.reposDBPath % 'postgres'),
                                      cfg.reposDBDriver)
                rcu = rdb.cursor()
                rcu.execute("SELECT datname FROM pg_database WHERE datname=?",
                             pgRepo)
                if rcu.fetchone():
                    rcu.execute("DROP DATABASE %s" % pgRepo)
                    rcu.execute("CREATE DATABASE %s ENCODING 'UTF8'" % pgRepo)
                rdb.close()

                dbUser = cfg.reposDBPath.split('@')[0]
                util.execute('cat %s | psql -U %s %s'% (dumpPath, dbUser, pgRepo))
                rdb = dbstore.connect('%s' % (cfg.reposDBPath % 'postgres'),
                                      cfg.reposDBDriver)
                conary.server.schema.loadSchema(rdb)
                rdb.close()
        else:
            cu.execute("SELECT * FROM InboundMirrors WHERE targetProjectId=?", projectId)
            localMirror = cu.fetchone_dict()
            if localMirror:
                if not os.path.exists(reposPath):
                    # revert Labels table to pre-mirror settings
                    cu.execute( \
                            "UPDATE Labels SET url=?, username=?, password=?" \
                                " WHERE projectId=?",
                        localMirror['sourceUrl'],
                        localMirror['sourceUsername'],
                        localMirror['sourcePassword'], projectId)

                    cu.execute("DELETE FROM RepNameMap WHERE toName=?", repo)
                    cu.execute( \
                        "DELETE FROM InboundMirrors WHERE inboundMirrorId=?",
                        localMirror['inboundMirrorId'])

    # delete package index to ensure we don't reference troves until they're
    # restored.
    try:
        cu.execute('DELETE FROM PackageIndex')
        cu.execute('DELETE FROM PackageIndexMark')
    except:
        # this masks errors, since it does re-raise but we don't care.
        # impact and cost are negligible.
        db.rollback()
    else:
        db.commit()

def prerestore(cfg):
    util.execute("service httpd stop")
    for repo in os.listdir(cfg.reposPath):
        reposPath = os.path.join(cfg.reposPath, repo)
        util.rmtree(reposPath, ignore_errors = True)

def clean(cfg):
    backupPath = os.path.join(cfg.dataPath, 'tmp', 'backup')
    util.rmtree(backupPath, ignore_errors = True)

def metadata(cfg, out):
    db = dbstore.connect(cfg.dbPath, cfg.dbDriver)
    cu = db.cursor()
    cu.execute("SELECT COALESCE(MAX(version), 0) FROM DatabaseVersion")
    schemaVersion = cu.fetchone()[0]
    print >> out, "rBuilder_schemaVersion=%d" % schemaVersion
    cu.execute("SELECT COALESCE(MAX(minor), 0) FROM DatabaseVersion")
    schemaMinor = cu.fetchone()[0]
    print >> out, "rBuilder_schemaMinor=%d" % schemaMinor
    sys.stdout.flush()

# XXX This version is used internally since our handler expects
#     an exception, while rAPA expects a True/False return.
def _isValid(cfg, input):
    if not isValid(cfg, input):
        raise RuntimeError("Invalid Backup")

def isValid(cfg, input):
    data = input.read()
    metaData = dict()
    for x in data.splitlines():
        if '=' in x:
           k, v = x.split('=', 1)
           metaData[k] = v
    schemaVersion = int(metaData.get('rBuilder_schemaVersion', 0))
    if schemaVersion == 0 or schemaVersion < schemaCutoff:
        return False
    NVF = metaData.get('NVF')
    if not NVF:
        return False
    verStr = cmdline.parseTroveSpec(NVF)[1]
    try:
        trailingVer = versions.VersionFromString(verStr).trailingRevision().version
    except conary.errors.ParseError:
        return False
    foundMatch = False
    for pat in knownGroupVersions:
        if re.match(pat, trailingVer):
            foundMatch = True
            break
    return foundMatch

    return True

def usage(out = sys.stderr):
    print >> out, sys.argv[0] + ":"
    print >> out, "    [b/backup]: back up databases and issue manifest."
    print >> out, "    [r/restore]: restore databases."
    print >> out, "    [c/clean]: clean up temporary files used during backup."
    print >> out, "    [m/metadata]: print appliance metadata in name=value format."
    print >> out, "    [isValid]: determine if metadata on stdin is for a compatible backup."
    print >> out, "    [prerestore]: prepare to restore databases."

def handle(func, *args, **kwargs):
    errno = 0
    if kwargs.pop('dropPriv', True):
        apacheUID, apacheGID = pwd.getpwnam('apache')[2:4]
        os.setgid(apacheGID)
        os.setuid(apacheUID)
    try:
        func(*args, **kwargs)
    except:
        exception, e, bt = sys.exc_info()
        if exception in (OSError, IOError):
            errno = e.errno
        else:
            errno = 1
        import traceback
        print >> sys.stderr, ''.join(traceback.format_tb(bt))
        print >> sys.stderr, exception, e
    sys.exit(errno)

def run():
    cfg = config.MintConfig()
    cfg.read(config.RBUILDER_CONFIG)
    migration = bool(os.getenv('RBA_MIGRATION'))
    mode = (len(sys.argv) > 1) and sys.argv[1] or ''
    if mode in ('r', 'restore'):
        handle(restore, cfg)
    elif mode  == 'prerestore':
        handle(prerestore, cfg, dropPriv = False)
    elif mode in ('b', 'backup'):
        handle(backup, cfg, sys.stdout, backupMirrors = migration)
    elif mode in ('c', 'clean'):
        handle(clean, cfg)
    elif mode in ('m', 'metadata'):
        handle(metadata, cfg, sys.stdout)
    elif mode == 'isValid':
        handle(_isValid, cfg, sys.stdin)
    elif mode.upper() in ('?', 'H', 'HELP'):
        usage(out = sys.stdout)
    else:
        usage()

<commit_msg>remove extraneous return<commit_after>#
# Copyright (c) 2005-2007 rPath, Inc.
#
# All rights reserved
#

import re, os, sys
import pwd
from mint import config
from mint import projects
from mint import schema
from conary.lib import util
from conary import dbstore
from conary import versions
import conary.errors
import conary.server.schema
from conary.conaryclient import cmdline

schemaCutoff = 37
knownGroupVersions = ('3\.1\.5.*', '4\..*')

staticPaths = ['config', 'entitlements', 'logs', 'installable_iso.conf',
               'toolkit', 'iso_gen.conf', 'live_iso.conf',
               'bootable_image.conf']
# toolkit is linked to outgoing job server architecture.

def backup(cfg, out, backupMirrors = False):
    reposContentsDir = cfg.reposContentsDir.split(' ')
    backupPath = os.path.join(cfg.dataPath, 'tmp', 'backup')
    print >> out, backupPath
    util.mkdirChain(backupPath)
    dumpPath = os.path.join(backupPath, 'db.dump')
    if cfg.dbDriver == 'sqlite':
        util.execute("echo '.dump' | sqlite3 %s > %s" % (cfg.dbPath, dumpPath))
    extraArgs = ""
    if not backupMirrors:
        extraArgs = " WHERE NOT external"
    if cfg.reposDBDriver == 'sqlite':
        db = dbstore.connect(cfg.dbPath, cfg.dbDriver)
        cu = db.cursor()
        cu.execute("SELECT hostname, domainname FROM Projects%s" % extraArgs)
        for reposDir in ['.'.join(x) for x in cu.fetchall()]:
            cu.execute("SELECT toName FROM RepNameMap WHERE fromName=?",
                    reposDir)
            res = cu.fetchone()
            reposDir = res and res[0] or reposDir
            reposDBPath = cfg.reposDBPath % reposDir
            dumpPath = os.path.join(backupPath, '%s.dump' % reposDir)
            if os.path.exists(reposDBPath):
                util.execute("echo '.dump' | sqlite3 %s > %s" % \
                              (reposDBPath, dumpPath))
                print >> out, reposContentsDir[0] % reposDir
    elif cfg.reposDBDriver == 'postgresql':
        db = dbstore.connect(cfg.dbPath, cfg.dbDriver)
        cu = db.cursor()
        cu.execute("SELECT hostname, domainname FROM Projects%s" % extraArgs)
        for reposDir in ['.'.join(x) for x in cu.fetchall()]:
            reposDbName = reposDir.translate(projects.transTables['postgresql'])
            dbUser = cfg.reposDBPath.split('@')[0]
            dumpPath = os.path.join(backupPath, '%s.dump' % reposDbName)
            rdb = dbstore.connect(cfg.reposDBPath % 'postgres', 'postgresql')
            rcu = rdb.cursor()
            rcu.execute("SELECT datname FROM pg_database WHERE datname=?",
                         reposDbName)
            if rcu.fetchone():
                util.execute("pg_dump -U %s -c -O -d %s > %s" %\
                             (dbUser, reposDbName, dumpPath))
                print >> out, reposContentsDir[0] % reposDir
            rdb.close()

    for d in staticPaths:
        path = os.path.join(cfg.dataPath, d)
        if os.path.exists(path):
            print >> out, path

def restore(cfg):
    backupPath = os.path.join(cfg.dataPath, 'tmp', 'backup')
    dumpPath = os.path.join(backupPath, 'db.dump')
    if cfg.dbDriver == 'sqlite':
        if os.path.exists(cfg.dbPath):
            os.unlink(cfg.dbPath)
        util.execute("sqlite3 %s < %s" % (cfg.dbPath, dumpPath))
    db = dbstore.connect(cfg.dbPath, cfg.dbDriver)
    schema.loadSchema(db, cfg)
    cu = db.cursor()
    cu.execute("SELECT hostname, domainname, projectId FROM Projects")
    for hostname, domainname, projectId in [x for x in cu.fetchall()]:
        repo = hostname + '.' + domainname
        cu.execute('SELECT toName FROM RepNameMap WHERE fromName=?', repo)
        r = cu.fetchone()
        if r:
            repo = r[0]
        dumpPath =  os.path.join(backupPath, repo).translate(projects.transTables[cfg.reposDBDriver]) + '.dump'
        reposPath = os.path.join(cfg.reposPath, repo)
        if os.path.exists(dumpPath):
            util.mkdirChain(os.path.join(reposPath, 'tmp'))
            if cfg.reposDBDriver == 'sqlite':
                dbPath = cfg.reposDBPath % repo
                if os.path.exists(dbPath):
                    os.unlink(dbPath)
                util.execute('sqlite3 %s < %s' % (dbPath, dumpPath))
                rdb = dbstore.connect(dbPath, cfg.reposDBDriver)
                conary.server.schema.loadSchema(rdb)
                rdb.close()
            elif cfg.reposDBDriver == 'postgresql':
                pgRepo = repo.translate(projects.transTables['postgresql'])
                dbPath = cfg.reposDBPath % pgRepo
                rdb = dbstore.connect('%s' % (cfg.reposDBPath % 'postgres'),
                                      cfg.reposDBDriver)
                rcu = rdb.cursor()
                rcu.execute("SELECT datname FROM pg_database WHERE datname=?",
                             pgRepo)
                if rcu.fetchone():
                    rcu.execute("DROP DATABASE %s" % pgRepo)
                    rcu.execute("CREATE DATABASE %s ENCODING 'UTF8'" % pgRepo)
                rdb.close()

                dbUser = cfg.reposDBPath.split('@')[0]
                util.execute('cat %s | psql -U %s %s'% (dumpPath, dbUser, pgRepo))
                rdb = dbstore.connect('%s' % (cfg.reposDBPath % 'postgres'),
                                      cfg.reposDBDriver)
                conary.server.schema.loadSchema(rdb)
                rdb.close()
        else:
            cu.execute("SELECT * FROM InboundMirrors WHERE targetProjectId=?", projectId)
            localMirror = cu.fetchone_dict()
            if localMirror:
                if not os.path.exists(reposPath):
                    # revert Labels table to pre-mirror settings
                    cu.execute( \
                            "UPDATE Labels SET url=?, username=?, password=?" \
                                " WHERE projectId=?",
                        localMirror['sourceUrl'],
                        localMirror['sourceUsername'],
                        localMirror['sourcePassword'], projectId)

                    cu.execute("DELETE FROM RepNameMap WHERE toName=?", repo)
                    cu.execute( \
                        "DELETE FROM InboundMirrors WHERE inboundMirrorId=?",
                        localMirror['inboundMirrorId'])

    # delete package index to ensure we don't reference troves until they're
    # restored.
    try:
        cu.execute('DELETE FROM PackageIndex')
        cu.execute('DELETE FROM PackageIndexMark')
    except:
        # this masks errors, since it does re-raise but we don't care.
        # impact and cost are negligible.
        db.rollback()
    else:
        db.commit()

def prerestore(cfg):
    util.execute("service httpd stop")
    for repo in os.listdir(cfg.reposPath):
        reposPath = os.path.join(cfg.reposPath, repo)
        util.rmtree(reposPath, ignore_errors = True)

def clean(cfg):
    backupPath = os.path.join(cfg.dataPath, 'tmp', 'backup')
    util.rmtree(backupPath, ignore_errors = True)

def metadata(cfg, out):
    db = dbstore.connect(cfg.dbPath, cfg.dbDriver)
    cu = db.cursor()
    cu.execute("SELECT COALESCE(MAX(version), 0) FROM DatabaseVersion")
    schemaVersion = cu.fetchone()[0]
    print >> out, "rBuilder_schemaVersion=%d" % schemaVersion
    cu.execute("SELECT COALESCE(MAX(minor), 0) FROM DatabaseVersion")
    schemaMinor = cu.fetchone()[0]
    print >> out, "rBuilder_schemaMinor=%d" % schemaMinor
    sys.stdout.flush()

# XXX This version is used internally since our handler expects
#     an exception, while rAPA expects a True/False return.
def _isValid(cfg, input):
    if not isValid(cfg, input):
        raise RuntimeError("Invalid Backup")

def isValid(cfg, input):
    data = input.read()
    metaData = dict()
    for x in data.splitlines():
        if '=' in x:
           k, v = x.split('=', 1)
           metaData[k] = v
    schemaVersion = int(metaData.get('rBuilder_schemaVersion', 0))
    if schemaVersion == 0 or schemaVersion < schemaCutoff:
        return False
    NVF = metaData.get('NVF')
    if not NVF:
        return False
    verStr = cmdline.parseTroveSpec(NVF)[1]
    try:
        trailingVer = versions.VersionFromString(verStr).trailingRevision().version
    except conary.errors.ParseError:
        return False
    foundMatch = False
    for pat in knownGroupVersions:
        if re.match(pat, trailingVer):
            foundMatch = True
            break
    return foundMatch

def usage(out = sys.stderr):
    print >> out, sys.argv[0] + ":"
    print >> out, "    [b/backup]: back up databases and issue manifest."
    print >> out, "    [r/restore]: restore databases."
    print >> out, "    [c/clean]: clean up temporary files used during backup."
    print >> out, "    [m/metadata]: print appliance metadata in name=value format."
    print >> out, "    [isValid]: determine if metadata on stdin is for a compatible backup."
    print >> out, "    [prerestore]: prepare to restore databases."

def handle(func, *args, **kwargs):
    errno = 0
    if kwargs.pop('dropPriv', True):
        apacheUID, apacheGID = pwd.getpwnam('apache')[2:4]
        os.setgid(apacheGID)
        os.setuid(apacheUID)
    try:
        func(*args, **kwargs)
    except:
        exception, e, bt = sys.exc_info()
        if exception in (OSError, IOError):
            errno = e.errno
        else:
            errno = 1
        import traceback
        print >> sys.stderr, ''.join(traceback.format_tb(bt))
        print >> sys.stderr, exception, e
    sys.exit(errno)

def run():
    cfg = config.MintConfig()
    cfg.read(config.RBUILDER_CONFIG)
    migration = bool(os.getenv('RBA_MIGRATION'))
    mode = (len(sys.argv) > 1) and sys.argv[1] or ''
    if mode in ('r', 'restore'):
        handle(restore, cfg)
    elif mode  == 'prerestore':
        handle(prerestore, cfg, dropPriv = False)
    elif mode in ('b', 'backup'):
        handle(backup, cfg, sys.stdout, backupMirrors = migration)
    elif mode in ('c', 'clean'):
        handle(clean, cfg)
    elif mode in ('m', 'metadata'):
        handle(metadata, cfg, sys.stdout)
    elif mode == 'isValid':
        handle(_isValid, cfg, sys.stdin)
    elif mode.upper() in ('?', 'H', 'HELP'):
        usage(out = sys.stdout)
    else:
        usage()

<|endoftext|>