<commit_before>from collections import defaultdict
from simuvex.s_ref import RefTypes, SimRegRead, SimRegWrite, SimTmpRead, SimTmpWrite, SimMemRead, SimMemWrite
from simuvex import SimIRSB, SimProcedure
import logging

l = logging.getLogger("angr.ddg")
l.setLevel(logging.DEBUG)

class DDG(object):
    def __init__(self, cfg, entry_point):
        self._cfg = cfg
        self._entry_point = entry_point

        self._ddg = defaultdict(lambda:defaultdict(set))
        self._symbolic_mem_ops = set()

    def debug_print(self):
        l.debug(self._ddg)

    def _trace_source(self, init_run, init_ref):
        '''
        Trace the sources (producers) of a given ref
        '''
        l.debug("Tracing source for symbolic memory dep %s of %s", init_ref, init_run)
        sources = set()
        # Memorization
        traced = {}

        init_stmt_id = init_ref.stmt_idx
        init_reg_deps = set()
        init_tmp_deps = set()
        real_ref = init_run.statements[init_stmt_id].refs[-1]
        if type(real_ref) == SimMemWrite:
            init_reg_deps |= set(list(real_ref.data_reg_deps))
            init_tmp_deps |= set(list(real_ref.data_tmp_deps))
            init_reg_deps |= set(list(real_ref.addr_reg_deps))
            init_tmp_deps |= set(list(real_ref.addr_tmp_deps))
        elif type(real_ref) == SimMemRead:
            init_reg_deps |= set(list(real_ref.addr_reg_deps))
            init_tmp_deps |= set(list(real_ref.addr_tmp_deps))
        else:
            init_reg_deps |= set(list(real_ref.data_reg_deps))
            init_tmp_deps |= set(list(real_ref.data_tmp_deps))
        stack = [(init_run, init_stmt_id, init_reg_deps, init_tmp_deps)]
        while len(stack) > 0:
            run, stmt_id, reg_deps, tmp_deps = stack.pop()
            l.debug("Traversing %s", run)
            # Added to traced set
            traced[run.addr] = reg_deps
            reg_dep_to_stmt_id = {}

            if isinstance(run, SimIRSB):
                irsb = run
                if stmt_id == -1:
                    stmt_ids = range(len(irsb.statements))
                else:
                    stmt_ids = range(stmt_id + 1)
                stmt_ids.reverse()

                for stmt_id in stmt_ids:
                    stmt = irsb.statements[stmt_id]
                    if len(stmt.refs) > 0:
                        real_ref = stmt.refs[-1]
                        if type(real_ref) == SimRegWrite and real_ref.offset in reg_deps:
                            reg_dep_to_stmt_id[real_ref.offset] = stmt_id
                            reg_deps.remove(real_ref.offset)
                            reg_deps |= set(list(real_ref.data_reg_deps))
                            tmp_deps |= set(list(real_ref.data_tmp_deps))
                        elif type(real_ref) == SimTmpWrite and real_ref.tmp in tmp_deps:
                            tmp_deps.remove(real_ref.tmp)
                            reg_deps |= set(list(real_ref.data_reg_deps))
                            tmp_deps |= set(list(real_ref.data_tmp_deps))
            elif isinstance(run, SimProcedure):
                refs = run.refs()
                refs.reverse()
                for ref in refs:
                    if type(ref) == SimRegWrite and ref.offset in reg_deps:
                        reg_dep_to_stmt_id[ref.offset] = -1
                        reg_deps.remove(ref.offset)
                        reg_deps |= set(list(ref.data_reg_deps))
                        tmp_deps |= set(list(ref.data_tmp_deps))
                    elif type(ref) == SimTmpWrite and ref.tmp in tmp_deps:
                        tmp_deps.remove(ref.tmp)
                        reg_deps |= set(list(ref.data_reg_deps))
                        tmp_deps |= set(list(ref.data_tmp_deps))

            if len(reg_deps) > 0:
                predecessors = self._cfg.get_predecessors(run)
                for s in predecessors:
                    tpl = (s.addr, reg_deps)
                    if s.addr not in traced:
                        stack.append((s, -1, reg_deps.copy(), set()))
                    else:
                        old_reg_deps = traced[s.addr]
                        if not reg_deps.issubset(old_reg_deps):
                            new_reg_deps = reg_deps.copy()
                            new_reg_deps |= old_reg_deps
                            stack.append((s, -1, new_reg_deps, set()))
            for reg, stmt_id in reg_dep_to_stmt_id.items():
                sources.add((run.addr, stmt_id))
            if len(reg_deps) == 0 or len(predecessors) == 0:
                # This is the end of the u-d chain
                # FIXME: What if there is a loop, and len(successors) never == 0?
                for reg in reg_deps:
                    if reg not in reg_dep_to_stmt_id:
                        # It's never been written before
                        # We use a fake run.addr and stmt_id so that we don't lose them
                        l.debug("Register %d has never been assigned a value before.", reg)
                        sources.add((-1 * reg, -1))
                    else:
                        # sources.add((run.addr, reg_dep_to_stmt_id[reg]))
                        # They have already been added to sources
                        pass

        return sources

    def _solve_symbolic_mem_operations(self):
        '''
        We try to resolve symbolic memory operations in the following manner:
        For each memory operation, trace the pointer from its root producers.
        And then relate each memory read with all memory writes that has shared
        producers with it.
        It's imprecise and could be over-approximating, but it's better than
        losing dependencies.
        '''
        # irsb, stmt_id => list of tuples (irsb, ref)
        mem_read_producers = defaultdict(list)
        mem_write_producers = defaultdict(list)
        for irsb, ref in self._symbolic_mem_ops:
            sources = self._trace_source(irsb, ref)
            for src_irsb_addr, stmt_id in sources:
                if type(ref) == SimMemRead:
                    mem_read_producers[(src_irsb_addr, stmt_id)].append((irsb.addr, ref))
                elif type(ref) == SimMemWrite:
                    mem_write_producers[(src_irsb_addr, stmt_id)].append((irsb.addr, ref))

        for tpl, lst_writes in mem_write_producers.items():
            if tpl in mem_read_producers:
                lst_reads = mem_read_producers[tpl]
                for read_irsb, read_ref in lst_reads:
                    for write_irsb, write_ref in lst_writes:
                        self._ddg[read_irsb][read_ref.stmt_idx].add((write_irsb, write_ref.stmt_idx))

    def construct(self):
        '''
    We track the following types of memory access:
    - (Intra-functional) Stack read/write.
        Trace changes of stack pointers inside a function, and the dereferences
        of stack pointers.
    - (Inter-functional) Stack read/write.
        TODO
    - (Global) Static memory positions.
        Keep a map of all accessible memory positions to their source statements
        per function. After that, we traverse the CFG and link each pair of
        reads/writes together in the order of control-flow.
    - (Intra-functional) Indirect memory access
        TODO

    In this way it's an O(N) algorithm, where N is the number of all functions
    (with context sensitivity). The old worklist algorithm is exponential in the
    worst case.
        '''
        # Stack access
        # Unfortunately, we have no idea where a function is, and it's better
        # not to rely on function identification methods. So we just traverse
        # the CFG once, and maintain a map of scanned IRSBs so that we scan
        # each IRSB only once.
        scanned_runs = set()
        initial_irsb = self._cfg.get_irsb((None, None, self._entry_point))
        # We maintain a calling stack so that we can limit all analysis within
        # the range of a single function.
        # Without function analysis, our approach is quite simple: If there is
        # an SimExit whose jumpkind is 'Ijk_Call', then we create a new frame
        # in calling stack. A frame is popped out if there exists an SimExit
        # that has 'Ijk_Ret' as its jumpkind.
        initial_wrapper = RunWrapper(initial_irsb)
        # All pending SimRuns
        run_stack = [initial_wrapper]
        while len(run_stack) > 0:
            run_wrapper = run_stack.pop()

            # Get the current calling stack
            current_stack = run_wrapper.call_stack

            run = run_wrapper.run
            if run in scanned_runs:
                continue
            scanned_runs.add(run)
            l.debug("Scanning %s", run)

            # Expand the current SimRun
            successors = self._cfg.get_successors(run)
            # TODO: It is buggy here. We are losing correlation between run.exits
            # and each successor in the CFG!
            for successor in successors:
                if successor in scanned_runs:
                    continue

                new_stack = current_stack[::] # Make a copy
                if run.exits()[0].jumpkind == "Ijk_Call":
                    # Create a new function frame
                    new_stack.append(run_wrapper)
                elif run.exits()[0].jumpkind == "Ijk_Ret":
                    if len(new_stack) > 0:
                        # Pop out the latest function frame
                        new_stack.pop()
                    else:
                        # We are returning from somewhere, but the stack is
                        # already empty.
                        # Something must have went wrong.
                        l.warning("Stack is already empty before popping things out")
                else:
                    # Do nothing :)
                    pass
                wrapper = RunWrapper(successor, \
                        addr_to_ref=run_wrapper.addr_to_ref, \
                        call_stack=new_stack)
                run_stack.append(wrapper)

            if isinstance(run, SimIRSB):
                irsb = run

                # Simulate the execution of this IRSB.
                # For MemWriteRef, fill the addr_to_ref dict with every single
                # concretizable memory address, and just ignore symbolic ones
                # for now.
                # For MemReadRef, we get its related MemWriteRef from our dict
                stmts = irsb.statements
                for i in range(len(stmts)):
                    stmt = stmts[i]
                    refs = stmt.refs
                    if len(refs) == 0:
                        continue
                    # Obtain the real Ref of current statement
                    real_ref = refs[-1]
                    if type(real_ref) == SimMemWrite:
                        addr = real_ref.addr
                        if not addr.is_symbolic():
                            # It's not symbolic. Try to concretize it.
                            concrete_addr = addr.any()
                            # Create the tuple of (simrun_addr, stmt_id)
                            tpl = (irsb.addr, i)
                            run_wrapper.addr_to_ref[concrete_addr] = tpl
                            l.debug("Memory write to addr 0x%x, stmt id = %d", concrete_addr, i)
                        else:
                            # Add it to our symbolic memory operation list. We
                            # will process them later.
                            self._symbolic_mem_ops.add((irsb, real_ref))
                    for ref in refs:
                        if type(ref) == SimMemRead:
                            addr = ref.addr
                            if not addr.is_symbolic():
                                # Not symbolic. Try to concretize it.
                                concrete_addr = addr.any()
                                # Check if this address has been written before.
                                # Note: we should check every single call frame,
                                # from the latest to earliest, until we come
                                # across that address.
                                reversed_range = range(len(current_stack))
                                reversed_range.reverse()
                                for j in reversed_range:
                                    con_ = current_stack[j]
                                    if concrete_addr in con_.addr_to_ref:
                                        # Luckily we found it!
                                        # Record it in our internal dict
                                        self._ddg[irsb.addr][i].add(
                                            con_.addr_to_ref[concrete_addr])
                                        break
                                # TODO: What if we have never seen this address
                                # before? It might be an address whose value is
                                # initialized somewhere else, or an address that
                                # contains initialized value.
                            else:
                                self._symbolic_mem_ops.add((irsb, ref))
            else:
                # SimProcedure
                sim_proc = run
                l.debug("Scanning %s", sim_proc)

                refs = sim_proc.refs()
                for ref in refs:
                    if isinstance(ref, SimMemWrite):
                        addr = ref.addr
                        if not addr.is_symbolic():
                            # Record it
                            # Not symbolic. Try to concretize it.
                            concrete_addr = addr.any()
                            # Create the tuple of (simrun_addr, stmt_id)
                            tpl = (sim_proc.addr, i)
                            run_wrapper.addr_to_ref[concrete_addr] = tpl
                        else:
                            self._symbolic_mem_ops.add((sim_proc, ref))
                    elif isinstance(ref, SimMemRead):
                        addr = ref.addr
                        if not addr.is_symbolic():
                            # Not symbolic. Try to concretize it.
                            concrete_addr = addr.any()
                            reversed_range = range(len(current_stack))
                            reversed_range.reverse()
                            for j in reversed_range:
                                con_ = current_stack[j]
                                if concrete_addr in con_.addr_to_ref:
                                    # Luckily we found it!
                                    # Record it in our internal dict
                                    self._ddg[sim_proc.addr][-1].add(
                                        con_.addr_to_ref[concrete_addr])
                                    break
                            # TODO: what if we didn't see that address before?
                        else:
                            self._symbolic_mem_ops.add((sim_proc, ref))

        self._solve_symbolic_mem_operations()

class RunWrapper(object):
    '''
We keep an RunWrapper object for each function (with context sensitivity, it's
an AddrToRefContainer object for each [function, context] pair). It contains a
list of all runs inside this function, a dict addr_to_ref storing all
references between addresses and a [simrun_addr, stmt_id] pair, and a calling
stack.
    '''
# TODO: We might want to change the calling stack into a branching list with
# CoW supported.
    def __init__(self, run, addr_to_ref=None, call_stack=None, reanalyze_successors=False):
        self.run = run
        if addr_to_ref is None:
            self.addr_to_ref = {}
        else:
            self.addr_to_ref = addr_to_ref

        if call_stack is None:
            self.call_stack = []
        else:
            # We DO NOT make a copy of the provided stack object
            self.call_stack = call_stack

        self.reanalyze_successors = reanalyze_successors
'''
    def combine(self, new_addr_to_ref):
        altered = False
        for k, v in new_addr_to_ref.items():
            if k not in self.addr_to_ref:
                self.addr_to_ref[k] = v.copy()
                altered = True
            else:
                if not v.issubset(self.addr_to_ref[k]):
                    altered = True
                    self.addr_to_ref[k] |= v
        return altered
'''
<commit_msg>Fixed a bug in DDG generation that may cause losing data dependences after branches.<commit_after>from collections import defaultdict
from simuvex.s_ref import RefTypes, SimRegRead, SimRegWrite, SimTmpRead, SimTmpWrite, SimMemRead, SimMemWrite
from simuvex import SimIRSB, SimProcedure
import logging

l = logging.getLogger("angr.ddg")
l.setLevel(logging.DEBUG)

class DDG(object):
    def __init__(self, cfg, entry_point):
        self._cfg = cfg
        self._entry_point = entry_point

        self._ddg = defaultdict(lambda:defaultdict(set))
        self._symbolic_mem_ops = set()

    def debug_print(self):
        l.debug(self._ddg)

    def _trace_source(self, init_run, init_ref):
        '''
        Trace the sources (producers) of a given ref
        '''
        l.debug("Tracing source for symbolic memory dep %s of %s", init_ref, init_run)
        sources = set()
        # Memorization
        traced = {}

        init_stmt_id = init_ref.stmt_idx
        init_reg_deps = set()
        init_tmp_deps = set()
        real_ref = init_run.statements[init_stmt_id].refs[-1]
        if type(real_ref) == SimMemWrite:
            init_reg_deps |= set(list(real_ref.data_reg_deps))
            init_tmp_deps |= set(list(real_ref.data_tmp_deps))
            init_reg_deps |= set(list(real_ref.addr_reg_deps))
            init_tmp_deps |= set(list(real_ref.addr_tmp_deps))
        elif type(real_ref) == SimMemRead:
            init_reg_deps |= set(list(real_ref.addr_reg_deps))
            init_tmp_deps |= set(list(real_ref.addr_tmp_deps))
        else:
            init_reg_deps |= set(list(real_ref.data_reg_deps))
            init_tmp_deps |= set(list(real_ref.data_tmp_deps))
        stack = [(init_run, init_stmt_id, init_reg_deps, init_tmp_deps)]
        while len(stack) > 0:
            run, stmt_id, reg_deps, tmp_deps = stack.pop()
            l.debug("Traversing %s", run)
            # Added to traced set
            traced[run.addr] = reg_deps
            reg_dep_to_stmt_id = {}

            if isinstance(run, SimIRSB):
                irsb = run
                if stmt_id == -1:
                    stmt_ids = range(len(irsb.statements))
                else:
                    stmt_ids = range(stmt_id + 1)
                stmt_ids.reverse()

                for stmt_id in stmt_ids:
                    stmt = irsb.statements[stmt_id]
                    if len(stmt.refs) > 0:
                        real_ref = stmt.refs[-1]
                        if type(real_ref) == SimRegWrite and real_ref.offset in reg_deps:
                            reg_dep_to_stmt_id[real_ref.offset] = stmt_id
                            reg_deps.remove(real_ref.offset)
                            reg_deps |= set(list(real_ref.data_reg_deps))
                            tmp_deps |= set(list(real_ref.data_tmp_deps))
                        elif type(real_ref) == SimTmpWrite and real_ref.tmp in tmp_deps:
                            tmp_deps.remove(real_ref.tmp)
                            reg_deps |= set(list(real_ref.data_reg_deps))
                            tmp_deps |= set(list(real_ref.data_tmp_deps))
            elif isinstance(run, SimProcedure):
                refs = run.refs()
                refs.reverse()
                for ref in refs:
                    if type(ref) == SimRegWrite and ref.offset in reg_deps:
                        reg_dep_to_stmt_id[ref.offset] = -1
                        reg_deps.remove(ref.offset)
                        reg_deps |= set(list(ref.data_reg_deps))
                        tmp_deps |= set(list(ref.data_tmp_deps))
                    elif type(ref) == SimTmpWrite and ref.tmp in tmp_deps:
                        tmp_deps.remove(ref.tmp)
                        reg_deps |= set(list(ref.data_reg_deps))
                        tmp_deps |= set(list(ref.data_tmp_deps))

            if len(reg_deps) > 0:
                predecessors = self._cfg.get_predecessors(run)
                for s in predecessors:
                    tpl = (s.addr, reg_deps)
                    if s.addr not in traced:
                        stack.append((s, -1, reg_deps.copy(), set()))
                    else:
                        old_reg_deps = traced[s.addr]
                        if not reg_deps.issubset(old_reg_deps):
                            new_reg_deps = reg_deps.copy()
                            new_reg_deps |= old_reg_deps
                            stack.append((s, -1, new_reg_deps, set()))
            for reg, stmt_id in reg_dep_to_stmt_id.items():
                sources.add((run.addr, stmt_id))
            if len(reg_deps) == 0 or len(predecessors) == 0:
                # This is the end of the u-d chain
                # FIXME: What if there is a loop, and len(successors) never == 0?
                for reg in reg_deps:
                    if reg not in reg_dep_to_stmt_id:
                        # It's never been written before
                        # We use a fake run.addr and stmt_id so that we don't lose them
                        l.debug("Register %d has never been assigned a value before.", reg)
                        sources.add((-1 * reg, -1))
                    else:
                        # sources.add((run.addr, reg_dep_to_stmt_id[reg]))
                        # They have already been added to sources
                        pass

        return sources

    def _solve_symbolic_mem_operations(self):
        '''
        We try to resolve symbolic memory operations in the following manner:
        For each memory operation, trace the pointer from its root producers.
        And then relate each memory read with all memory writes that has shared
        producers with it.
        It's imprecise and could be over-approximating, but it's better than
        losing dependencies.
        '''
        # irsb, stmt_id => list of tuples (irsb, ref)
        mem_read_producers = defaultdict(list)
        mem_write_producers = defaultdict(list)
        for irsb, ref in self._symbolic_mem_ops:
            sources = self._trace_source(irsb, ref)
            for src_irsb_addr, stmt_id in sources:
                if type(ref) == SimMemRead:
                    mem_read_producers[(src_irsb_addr, stmt_id)].append((irsb.addr, ref))
                elif type(ref) == SimMemWrite:
                    mem_write_producers[(src_irsb_addr, stmt_id)].append((irsb.addr, ref))

        for tpl, lst_writes in mem_write_producers.items():
            if tpl in mem_read_producers:
                lst_reads = mem_read_producers[tpl]
                for read_irsb, read_ref in lst_reads:
                    for write_irsb, write_ref in lst_writes:
                        self._ddg[read_irsb][read_ref.stmt_idx].add((write_irsb, write_ref.stmt_idx))

    def construct(self):
        '''
    We track the following types of memory access:
    - (Intra-functional) Stack read/write.
        Trace changes of stack pointers inside a function, and the dereferences
        of stack pointers.
    - (Inter-functional) Stack read/write.
        TODO
    - (Global) Static memory positions.
        Keep a map of all accessible memory positions to their source statements
        per function. After that, we traverse the CFG and link each pair of
        reads/writes together in the order of control-flow.
    - (Intra-functional) Indirect memory access
        TODO

    In this way it's an O(N) algorithm, where N is the number of all functions
    (with context sensitivity). The old worklist algorithm is exponential in the
    worst case.
        '''
        # Stack access
        # Unfortunately, we have no idea where a function is, and it's better
        # not to rely on function identification methods. So we just traverse
        # the CFG once, and maintain a map of scanned IRSBs so that we scan
        # each IRSB only once.
        scanned_runs = set()
        initial_irsb = self._cfg.get_irsb((None, None, self._entry_point))
        # We maintain a calling stack so that we can limit all analysis within
        # the range of a single function.
        # Without function analysis, our approach is quite simple: If there is
        # an SimExit whose jumpkind is 'Ijk_Call', then we create a new frame
        # in calling stack. A frame is popped out if there exists an SimExit
        # that has 'Ijk_Ret' as its jumpkind.
        initial_wrapper = RunWrapper(initial_irsb)
        # All pending SimRuns
        run_stack = [initial_wrapper]
        while len(run_stack) > 0:
            run_wrapper = run_stack.pop()

            # Get the current calling stack
            current_stack = run_wrapper.call_stack

            run = run_wrapper.run
            if run in scanned_runs:
                continue
            scanned_runs.add(run)
            l.debug("Scanning %s", run)

            reanalyze_successors_flag = run_wrapper.reanalyze_successors

            if isinstance(run, SimIRSB):
                irsb = run

                # Simulate the execution of this IRSB.
                # For MemWriteRef, fill the addr_to_ref dict with every single
                # concretizable memory address, and just ignore symbolic ones
                # for now.
                # For MemReadRef, we get its related MemWriteRef from our dict
                stmts = irsb.statements
                for i in range(len(stmts)):
                    stmt = stmts[i]
                    refs = stmt.refs
                    if len(refs) == 0:
                        continue
                    # Obtain the real Ref of current statement
                    real_ref = refs[-1]
                    if type(real_ref) == SimMemWrite:
                        addr = real_ref.addr
                        if not addr.is_symbolic():
                            # It's not symbolic. Try to concretize it.
                            concrete_addr = addr.any()
                            # Create the tuple of (simrun_addr, stmt_id)
                            tpl = (irsb.addr, i)
                            run_wrapper.addr_to_ref[concrete_addr] = tpl
                            l.debug("Memory write to addr 0x%x, irsb %s, stmt id = %d", concrete_addr, irsb, i)
                            reanalyze_successors_flag = True
                        else:
                            # Add it to our symbolic memory operation list. We
                            # will process them later.
                            self._symbolic_mem_ops.add((irsb, real_ref))
                    for ref in refs:
                        if type(ref) == SimMemRead:
                            addr = ref.addr
                            if not addr.is_symbolic():
                                # Not symbolic. Try to concretize it.
                                concrete_addr = addr.any()
                                # Check if this address has been written before.
                                # Note: we should check every single call frame,
                                # from the latest to earliest, until we come
                                # across that address.
                                reversed_range = range(len(current_stack))
                                reversed_range.reverse()
                                for j in reversed_range:
                                    con_ = current_stack[j]
                                    if concrete_addr in con_.addr_to_ref:
                                        # Luckily we found it!
                                        # Record it in our internal dict
                                        self._ddg[irsb.addr][i].add(
                                            con_.addr_to_ref[concrete_addr])
                                        break
                                # TODO: What if we have never seen this address
                                # before? It might be an address whose value is
                                # initialized somewhere else, or an address that
                                # contains initialized value.
                            else:
                                self._symbolic_mem_ops.add((irsb, ref))
            else:
                # SimProcedure
                sim_proc = run
                l.debug("Scanning %s", sim_proc)

                refs = sim_proc.refs()
                for ref in refs:
                    if isinstance(ref, SimMemWrite):
                        addr = ref.addr
                        if not addr.is_symbolic():
                            # Record it
                            # Not symbolic. Try to concretize it.
                            concrete_addr = addr.any()
                            # Create the tuple of (simrun_addr, stmt_id)
                            tpl = (sim_proc.addr, i)
                            run_wrapper.addr_to_ref[concrete_addr] = tpl
                            reanalyze_successors_flag = True
                        else:
                            self._symbolic_mem_ops.add((sim_proc, ref))
                    elif isinstance(ref, SimMemRead):
                        addr = ref.addr
                        if not addr.is_symbolic():
                            # Not symbolic. Try to concretize it.
                            concrete_addr = addr.any()
                            reversed_range = range(len(current_stack))
                            reversed_range.reverse()
                            for j in reversed_range:
                                con_ = current_stack[j]
                                if concrete_addr in con_.addr_to_ref:
                                    # Luckily we found it!
                                    # Record it in our internal dict
                                    self._ddg[sim_proc.addr][-1].add(
                                        con_.addr_to_ref[concrete_addr])
                                    break
                            # TODO: what if we didn't see that address before?
                        else:
                            self._symbolic_mem_ops.add((sim_proc, ref))

            # Expand the current SimRun
            successors = self._cfg.get_successors(run)
            # TODO: It is buggy here. We are losing correlation between run.exits
            # and each successor in the CFG!
            for successor in successors:
                if successor in scanned_runs:
                    if reanalyze_successors_flag:
                        scanned_runs.remove(successor)
                    else:
                        continue

                new_stack = current_stack[::] # Make a copy
                if run.exits()[0].jumpkind == "Ijk_Call":
                    # Create a new function frame
                    new_stack.append(run_wrapper)
                elif run.exits()[0].jumpkind == "Ijk_Ret":
                    if len(new_stack) > 0:
                        # Pop out the latest function frame
                        new_stack.pop()
                    else:
                        # We are returning from somewhere, but the stack is
                        # already empty.
                        # Something must have went wrong.
                        l.warning("Stack is already empty before popping things out")
                else:
                    # Do nothing :)
                    pass
                wrapper = RunWrapper(successor, \
                        addr_to_ref=run_wrapper.addr_to_ref, \
                        call_stack=new_stack,
                        reanalyze_successors=reanalyze_successors_flag)
                run_stack.append(wrapper)

        self._solve_symbolic_mem_operations()

class RunWrapper(object):
    '''
We keep an RunWrapper object for each function (with context sensitivity, it's
an AddrToRefContainer object for each [function, context] pair). It contains a
list of all runs inside this function, a dict addr_to_ref storing all
references between addresses and a [simrun_addr, stmt_id] pair, and a calling
stack.
    '''
# TODO: We might want to change the calling stack into a branching list with
# CoW supported.
    def __init__(self, run, addr_to_ref=None, call_stack=None, reanalyze_successors=False):
        self.run = run
        if addr_to_ref is None:
            self.addr_to_ref = {}
        else:
            self.addr_to_ref = addr_to_ref

        if call_stack is None:
            self.call_stack = []
        else:
            # We DO NOT make a copy of the provided stack object
            self.call_stack = call_stack

        self.reanalyze_successors = reanalyze_successors
'''
    def combine(self, new_addr_to_ref):
        altered = False
        for k, v in new_addr_to_ref.items():
            if k not in self.addr_to_ref:
                self.addr_to_ref[k] = v.copy()
                altered = True
            else:
                if not v.issubset(self.addr_to_ref[k]):
                    altered = True
                    self.addr_to_ref[k] |= v
        return altered
'''
<|endoftext|>


<commit_before>import logging
import sys
from disco.ddfs import DDFS

from inferno.lib.lazy_property import lazy_property

log = logging.getLogger(__name__)

class Archiver(object):

    def __init__(self,
                 ddfs,
                 tags,
                 urls=None,
                 archive_prefix='processed',
                 archive_mode=False,
                 nuke_mode=False,
                 max_blobs=sys.maxint,
                 newest_first=True):
        self.tags = tags
        self.ddfs = ddfs
        self.max_blobs = max_blobs
        self.archive_mode = archive_mode
        self.nuke_mode = nuke_mode
        self.archive_prefix = archive_prefix
        self.newest_first = newest_first
        self.tag_map = self._build_tag_map(tags)
        self.urls = urls if urls is not None else []

    @property
    def blob_count(self):
        return len(self.job_blobs)

    @lazy_property
    def job_blobs(self):
        # job_blobs consist of blobs of tags and urls
        job_blobs = []
        for blobs in self.tag_map.itervalues():
            job_blobs += blobs
        job_blobs += self.urls
        return job_blobs

    def archive(self):
        if self.tag_map:
            try:
                self._archive_tags()
            except Exception as e:
                log.error('Archiving error: %s', e, exc_info=sys.exc_info())

    def nuke(self):
        try:
            self.ddfs.delete(self.tags)
            log.info('Deleted %s .', self.tags)
        except Exception as e:
                log.error('Deleting error: %s', e, exc_info=sys.exc_info())

    def _archive_tags(self):
        for tag, blobs in self.tag_map.iteritems():
            if tag.startswith(self.archive_prefix):
                log.info('Tag already starts with archive prefix: %s', tag)
            else:
                archive_name = self._get_archive_name(tag)
                self.ddfs.tag(archive_name, blobs)
                log.info('Archived %d blobs to %s', len(blobs), archive_name)

    def _build_tag_map(self, tags):
        tag_map = {}
        blob_count = 0

        source_tags, archived_blobs = self._source_and_archived_sets(tags)
        for tag in source_tags:
            for blob in self.ddfs.blobs(tag):
                #normalized_blob = tuple(sorted(blob))
                blob_name = self.get_blob_name(blob[0])
                if blob_name not in archived_blobs:
                    incoming_blobs = tag_map.setdefault(tag, [])
                    incoming_blobs.append(blob)
                    blob_count += 1

                if blob_count >= self.max_blobs:
                    return tag_map
        return tag_map

    def _source_and_archived_sets(self, tags):
        source_tags = set()
        archived_blobs = set()
        for prefix in tags:
            for tag in self.ddfs.list(prefix):
                source_tags.add(tag)
                archive_tag = self._get_archive_name(tag)
                if self.archive_mode and not tag.startswith(self.archive_prefix):
                    archived_blobs.update(self._normalized_blobs(archive_tag))
        source_tags = sorted(source_tags, reverse=self.newest_first)
        return source_tags, archived_blobs

    def get_blob_name(self, blob):
        return blob.rsplit('/', 1)[1]

    def _normalized_blobs(self, tag):
        rval = set()
        for blob in self.ddfs.blobs(tag):
            rval.add(self.get_blob_name(blob[0]))
        return rval

    def _get_archive_name(self, tag):
        tag_parts = tag.split(':')
        tag_parts[0] = self.archive_prefix
        return ':'.join(tag_parts)
<commit_msg>Assert that the number of reported blobs is not more than max_blobs.<commit_after>import logging
import sys
from disco.ddfs import DDFS

from inferno.lib.lazy_property import lazy_property

log = logging.getLogger(__name__)

class Archiver(object):

    def __init__(self,
                 ddfs,
                 tags,
                 urls=None,
                 archive_prefix='processed',
                 archive_mode=False,
                 nuke_mode=False,
                 max_blobs=sys.maxint,
                 newest_first=True):
        self.tags = tags
        self.ddfs = ddfs
        self.max_blobs = max_blobs
        self.archive_mode = archive_mode
        self.nuke_mode = nuke_mode
        self.archive_prefix = archive_prefix
        self.newest_first = newest_first
        self.tag_map = self._build_tag_map(tags)
        self.urls = urls if urls is not None else []

    @property
    def blob_count(self):
        return len(self.job_blobs)

    @lazy_property
    def job_blobs(self):
        # job_blobs consist of blobs of tags and urls
        job_blobs = []
        for blobs in self.tag_map.itervalues():
            job_blobs += blobs
        job_blobs += self.urls
        return job_blobs

    def archive(self):
        if self.tag_map:
            try:
                self._archive_tags()
            except Exception as e:
                log.error('Archiving error: %s', e, exc_info=sys.exc_info())

    def nuke(self):
        try:
            self.ddfs.delete(self.tags)
            log.info('Deleted %s .', self.tags)
        except Exception as e:
                log.error('Deleting error: %s', e, exc_info=sys.exc_info())

    def _archive_tags(self):
        for tag, blobs in self.tag_map.iteritems():
            if tag.startswith(self.archive_prefix):
                log.info('Tag already starts with archive prefix: %s', tag)
            else:
                archive_name = self._get_archive_name(tag)
                self.ddfs.tag(archive_name, blobs)
                log.info('Archived %d blobs to %s', len(blobs), archive_name)

    def _build_tag_map(self, tags):
        tag_map = {}
        blob_count = 0

        source_tags, archived_blobs = self._source_and_archived_sets(tags)
        for tag in source_tags:
            for blob in self.ddfs.blobs(tag):
                #normalized_blob = tuple(sorted(blob))
                blob_name = self.get_blob_name(blob[0])
                if blob_name not in archived_blobs:
                    incoming_blobs = tag_map.setdefault(tag, [])
                    if blob_count == self.max_blobs:
                        return tag_map
                    incoming_blobs.append(blob)
                    blob_count += 1

                assert blob_count <= self.max_blobs
                if blob_count == self.max_blobs:
                    return tag_map
        return tag_map

    def _source_and_archived_sets(self, tags):
        source_tags = set()
        archived_blobs = set()
        for prefix in tags:
            for tag in self.ddfs.list(prefix):
                source_tags.add(tag)
                archive_tag = self._get_archive_name(tag)
                if self.archive_mode and not tag.startswith(self.archive_prefix):
                    archived_blobs.update(self._normalized_blobs(archive_tag))
        source_tags = sorted(source_tags, reverse=self.newest_first)
        return source_tags, archived_blobs

    def get_blob_name(self, blob):
        return blob.rsplit('/', 1)[1]

    def _normalized_blobs(self, tag):
        rval = set()
        for blob in self.ddfs.blobs(tag):
            rval.add(self.get_blob_name(blob[0]))
        return rval

    def _get_archive_name(self, tag):
        tag_parts = tag.split(':')
        tag_parts[0] = self.archive_prefix
        return ':'.join(tag_parts)
<|endoftext|>


<commit_before>
import os
import time
import socket
import xml.etree.ElementTree as ET
from aquilon.server.dbwrappers.service import get_service
from twisted.python import log
from aquilon.server.processes import write_file

CCM_NOTIF = 1
CDB_NOTIF = 2

NOTIFICATION_TYPES = {
        CCM_NOTIF : "ccm",
        CDB_NOTIF : "cdb"
}

def build_index(config, session, profilesdir, clientNotify=True):
    '''
    Create an index of what profiles are available

    Compare the mtimes of everything in profiledir against
    an index file (profiles-info.xml). Produce a new index
    and send out notifications to "server modules" (as defined
    within the broker configuration). If clientNotify
    is True, then individual notifications are also sent
    to each host. If clientNotify is False, then the server modules
    will still be notified, but there is no processing of the
    individual hosts. Note that the broker has a config option
    send_notifications, which if false will turn off notifications
    unconditionally. Only if the broker config allows will the
    clientNotify be checked.

    '''
    profile_index = "profiles-info.xml"

    old_object_index = {}
    index_path = os.path.join(profilesdir, profile_index)
    if os.path.exists(index_path):
        try:
            tree = ET.parse(index_path)
            for profile in tree.getiterator("profile"):
                if (profile.text and profile.attrib.has_key("mtime")):
                    object = profile.text.strip()
                    if object:
                        if object.endswith(".xml"):
                            object = object[:-4]
                        old_object_index[object] = int(profile.attrib["mtime"])
        except Exception, e:
            log.msg("Error processing %s, continuing: %s" % (index_path, e))

    object_index = {}
    modified_index = {}
    for profile in os.listdir(profilesdir):
        if profile == profile_index:
            continue
        if not profile.endswith(".xml"):
            continue
        object = profile[:-4]
        object_index[object] = os.path.getmtime(
                os.path.join(profilesdir, profile))
        if (old_object_index.has_key(object)
                and object_index[object] > old_object_index[object]):
            #log.msg("xml for %s has a newer mtime, will notify" % host)
            modified_index[object] = object_index[object]

    content = []
    content.append("<?xml version='1.0' encoding='utf-8'?>")
    content.append("<profiles>")
    for object, mtime in object_index.items():
        content.append("<profile mtime='%d'>%s.xml</profile>"
                % (mtime, object))
    content.append("</profiles>")

    write_file(index_path, "\n".join(content))

    if config.has_option("broker", "server_notifications"):
        service_modules = {}
        for service in config.get("broker", "server_notifications").split():
            if service.strip():
                try:
                    # service may be unknown
                    srvinfo = get_service(session, service)
                    for instance in srvinfo.instances:
                        for sis in instance.servers:
                            service_modules[sis.system.fqdn] = 1
                except Exception, e:
                    log.msg("failed to lookup up server module %s: %s" % (service, e))
        send_notification(CDB_NOTIF, service_modules.keys())

    if (config.has_option("broker", "client_notifications")
        and config.getboolean("broker", "client_notifications")
        and clientNotify):
        send_notification(CCM_NOTIF, modified_index.keys())

def send_notification(ntype, machines):
    '''send CDP notification messages to a list of hosts. This
    are sent synchronously, but we don't wait (or care) for any
    reply, so it shouldn't be a problem.
    type should be CCM_NOTIF or CDB_NOTIF
    '''
    packet = NOTIFICATION_TYPES[ntype] + "\0" + str(int(time.time()))
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    try:
        port = socket.getservbyname("cdp")
    except:
        port = 7777

    for host in machines:
        try:
            ip = socket.gethostbyname(host)
            log.msg("Sending %s notification to %s"
                    % (NOTIFICATION_TYPES[ntype], host))
            sock.sendto(packet, (ip, port))

        except socket.gaierror:
            # This host is not known (yet), so we silently
            # discard the notification.
            pass

        except Exception, e:
            log.msg("Error notifying %s: %s" % (host, e))

<commit_msg>notifications/indexing of profiles now supports namespaces<commit_after>
import os
import time
import socket
import xml.etree.ElementTree as ET
from aquilon.server.dbwrappers.service import get_service
from twisted.python import log
from aquilon.server.processes import write_file

CCM_NOTIF = 1
CDB_NOTIF = 2

NOTIFICATION_TYPES = {
        CCM_NOTIF : "ccm",
        CDB_NOTIF : "cdb"
}

try:
    CDPPORT = socket.getservbyname("cdp")
except:
    CDPPORT = 7777

def build_index(config, session, profilesdir, clientNotify=True):
    '''
    Create an index of what profiles are available

    Compare the mtimes of everything in profiledir against
    an index file (profiles-info.xml). Produce a new index
    and send out notifications to "server modules" (as defined
    within the broker configuration). If clientNotify
    is True, then individual notifications are also sent
    to each host. If clientNotify is False, then the server modules
    will still be notified, but there is no processing of the
    individual hosts. Note that the broker has a config option
    send_notifications, which if false will turn off notifications
    unconditionally. Only if the broker config allows will the
    clientNotify be checked.

    '''
    profile_index = "profiles-info.xml"

    old_object_index = {}
    index_path = os.path.join(profilesdir, profile_index)
    if os.path.exists(index_path):
        try:
            tree = ET.parse(index_path)
            for profile in tree.getiterator("profile"):
                if (profile.text and profile.attrib.has_key("mtime")):
                    obj = profile.text.strip()
                    if obj:
                        if obj.endswith(".xml"):
                            obj = obj[:-4]
                        old_object_index[obj] = int(profile.attrib["mtime"])
        except Exception, e:
            log.msg("Error processing %s, continuing: %s" % (index_path, e))

    # object_index ties namespaced files to mtime
    object_index = {}
    # modified_index stores the subset of namespaced names that
    # have changed since the last index. The values are unused.
    modified_index = {}

    for root, _dirs, files in os.walk(profilesdir):
        for profile in files:
            if profile == profile_index:
                continue
            if not profile.endswith(".xml"):
                continue
            obj = os.path.join(root, profile[:-4])
            # Remove the common prefix: our profilesdir, so that the
            # remaining object name is relative to that root (+1 in order
            # to remove the slash separator)
            obj = obj[len(profilesdir)+1:]
            object_index[obj] = os.path.getmtime(
                    os.path.join(root, profile))
            if (old_object_index.has_key(obj) and
                object_index[obj] > old_object_index[obj]):
                modified_index[obj] = object_index[obj]

    content = []
    content.append("<?xml version='1.0' encoding='utf-8'?>")
    content.append("<profiles>")
    for obj, mtime in object_index.items():
        content.append("<profile mtime='%d'>%s.xml</profile>"
                % (mtime, obj))
    content.append("</profiles>")

    write_file(index_path, "\n".join(content))

    if config.has_option("broker", "server_notifications"):
        service_modules = {}
        for service in config.get("broker", "server_notifications").split():
            if service.strip():
                try:
                    # service may be unknown
                    srvinfo = get_service(session, service)
                    for instance in srvinfo.instances:
                        for sis in instance.servers:
                            service_modules[sis.system.fqdn] = 1
                except Exception, e:
                    log.msg("failed to lookup up server module %s: %s" % (service, e))
        count = send_notification(CDB_NOTIF, service_modules.keys())
        log.msg("sent %d server notifications" % count)

    if (config.has_option("broker", "client_notifications")
        and config.getboolean("broker", "client_notifications")
        and clientNotify):
        count = send_notification(CCM_NOTIF, modified_index.keys())
        log.msg("sent %d client notifications" % count)

def send_notification(ntype, modified):
    '''send CDP notification messages to a list of hosts.

    This are sent synchronously, but we don't wait (or care) for any
    reply, so it shouldn't be a problem. type should be CCM_NOTIF or CDB_NOTIF.
    'modified' is a dict of object names that may be namespaced. Object
    names that cannot be looked up in DNS are silently ignored.
    Returns the number of notifications that were sent.
    '''

    success = 0
    for obj in modified:
        # We need to clean the name, since it might
        # be namespaced. This (in effect) globalizes
        # all names. Perhaps we might want to do some
        # checks based on the namespace. Not for now.
        (_ns, _sep, host) = obj.rpartition('/')

        try:
            ip = socket.gethostbyname(host)
            packet = NOTIFICATION_TYPES[ntype] + "\0" + str(int(time.time()))
            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            sock.sendto(packet, (ip, CDPPORT))
            success = success + 1

        except socket.gaierror:
            # This hostname is unknown, so we silently
            # discard the notification.
            pass

        except Exception, e:
            log.msg("Error notifying %s: %s" % (host, e))

    return success
<|endoftext|>


<commit_before>import argparse
import base64
import json

import numpy as np
import socketio
import eventlet
import eventlet.wsgi
import time
from PIL import Image
from PIL import ImageOps
from flask import Flask, render_template
from io import BytesIO

from keras.models import load_model
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array

# Fix error with Keras and TensorFlow
import tensorflow as tf
tf.python.control_flow_ops = tf

sio = socketio.Server()
app = Flask(__name__)
model = None
prev_image_array = None

@sio.on('telemetry')
def telemetry(sid, data):
    # The current steering angle of the car
    steering_angle = data["steering_angle"]
    # The current throttle of the car
    throttle = data["throttle"]
    # The current speed of the car
    speed = data["speed"]
    # The current image from the center camera of the car
    imgString = data["image"]
    image = Image.open(BytesIO(base64.b64decode(imgString)))
    image = image.resize((image.width//4, image.height//4), Image.BILINEAR)
    image_array = np.asarray(image)
    transformed_image_array = image_array[None, :, :, :]
    # This model currently assumes that the features of the model are just the images. Feel free to change this.
    steering_angle = float(model.predict(transformed_image_array, batch_size=1))
    # The driving model currently just outputs a constant throttle. Feel free to edit this.
    throttle = 0.2
    print(steering_angle, throttle)
    send_control(steering_angle, throttle)

@sio.on('connect')
def connect(sid, environ):
    print("connect ", sid)
    send_control(0, 0)

def send_control(steering_angle, throttle):
    sio.emit("steer", data={
    'steering_angle': steering_angle.__str__(),
    'throttle': throttle.__str__()
    }, skip_sid=True)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Remote Driving')
    parser.add_argument('model', type=str,
    help='Path to model definition json. Model weights should be on the same path.')
    args = parser.parse_args()
    #with open(args.model, 'r') as jfile:
        # NOTE: if you saved the file by calling json.dump(model.to_json(), ...)
        # then you will have to call:
        #
        #   model = model_from_json(json.loads(jfile.read()))\
        #
        # instead.
        #model = model_from_json(jfile.read())

    #model.compile("adam", "mse")
    #weights_file = args.model.replace('json', 'h5')
    #model.load_weights(weights_file)
    model = load_model(args.model)

    # wrap Flask application with engineio's middleware
    app = socketio.Middleware(sio, app)

    # deploy as an eventlet WSGI server
    eventlet.wsgi.server(eventlet.listen(('', 4567)), app)

<commit_msg>Code cleanup and comments<commit_after>import argparse
import base64
import json

import numpy as np
import socketio
import eventlet
import eventlet.wsgi
import time
from PIL import Image
from PIL import ImageOps
from flask import Flask, render_template
from io import BytesIO

from keras.models import load_model
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array

# Fix error with Keras and TensorFlow
import tensorflow as tf
tf.python.control_flow_ops = tf

sio = socketio.Server()
app = Flask(__name__)
model = None
prev_image_array = None

@sio.on('telemetry')
def telemetry(sid, data):
    # The current steering angle of the car
    steering_angle = data["steering_angle"]
    # The current throttle of the car
    throttle = data["throttle"]
    # The current speed of the car
    speed = data["speed"]
    # The current image from the center camera of the car
    imgString = data["image"]
    image = Image.open(BytesIO(base64.b64decode(imgString)))
    # Scale the image similarly to the training
    image = image.resize((image.width//4, image.height//4), Image.BILINEAR)
    image_array = np.asarray(image)
    transformed_image_array = image_array[None, :, :, :]
    # This model currently assumes that the features of the model are just the images. Feel free to change this.
    steering_angle = float(model.predict(transformed_image_array, batch_size=1))
    # The driving model currently just outputs a constant throttle. Feel free to edit this.
    throttle = 0.2
    print(steering_angle, throttle)
    send_control(steering_angle, throttle)

@sio.on('connect')
def connect(sid, environ):
    print("connect ", sid)
    send_control(0, 0)

def send_control(steering_angle, throttle):
    sio.emit("steer", data={
    'steering_angle': steering_angle.__str__(),
    'throttle': throttle.__str__()
    }, skip_sid=True)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Remote Driving')
    parser.add_argument('model', type=str,
    help='Path to model definition json. Model weights should be on the same path.')
    args = parser.parse_args()
    # Simplified loading of the model from a single .h5 file
    model = load_model(args.model)

    # wrap Flask application with engineio's middleware
    app = socketio.Middleware(sio, app)

    # deploy as an eventlet WSGI server
    eventlet.wsgi.server(eventlet.listen(('', 4567)), app)

<|endoftext|>


<commit_before>__author__ = 'martin'

class MyQueue:
    """Will take items and place them in to a list so that a FIFO structure is used."""
    def __init__(self):
        self._items = []

    def __str__(self):
        output = ""
        for x in self._items:
            output = output + str(x) + " "
        return output

    def enqueue(self, item):
        self._items.insert(0, item)

    def dequeue(self):
        if MyQueue.isEmpty(self):
            return "The Stack is empty."
        else:
            return self._items.pop()

    def __len__(self):
        return len(self._items)

    def isEmpty(self):
        return not len(self._items)

    def peek(self):
        return self._items[-1]

class Stack:
    """Defines a Stack ADT with operations: Stack, isEmpty,
        push, pop, peek, and len."""
    def __init__(self):
        self._items = []

    def __iter__(self):
        return iter(self._items)

    def __len__(self):
        return len(self._items)

    def __str__(self):
        output = ""
        for x in self._items:
            output = str(x) + " " + output
        return "[ " + output + "]"

    def isEmpty(self):
        return not len(self._items)

    def push(self, item):
        self._items.append(item)

    def pop(self):
        if Stack.isEmpty(self):
            return "The Stack is empty."
        else:
            return self._items.pop()

    def peek(self):
        return self._items[-1]

    def next(self):
        if len(self._items) > 0:
            return self._items.pop()
        else:
            raise StopIteration<commit_msg>change author<commit_after>__author__ = 'Lothilius'

class MyQueue:
    """Will take items and place them in to a list so that a FIFO structure is used."""
    def __init__(self):
        self._items = []

    def __str__(self):
        output = ""
        for x in self._items:
            output = output + str(x) + " "
        return output

    def enqueue(self, item):
        self._items.insert(0, item)

    def dequeue(self):
        if MyQueue.isEmpty(self):
            return "The Stack is empty."
        else:
            return self._items.pop()

    def __len__(self):
        return len(self._items)

    def isEmpty(self):
        return not len(self._items)

    def peek(self):
        return self._items[-1]

class Stack:
    """Defines a Stack ADT with operations: Stack, isEmpty,
        push, pop, peek, and len."""
    def __init__(self):
        self._items = []

    def __iter__(self):
        return iter(self._items)

    def __len__(self):
        return len(self._items)

    def __str__(self):
        output = ""
        for x in self._items:
            output = str(x) + " " + output
        return "[ " + output + "]"

    def isEmpty(self):
        return not len(self._items)

    def push(self, item):
        self._items.append(item)

    def pop(self):
        if Stack.isEmpty(self):
            return "The Stack is empty."
        else:
            return self._items.pop()

    def peek(self):
        return self._items[-1]

    def next(self):
        if len(self._items) > 0:
            return self._items.pop()
        else:
            raise StopIteration<|endoftext|>


<commit_before>from cantools.web import respond, success, fail, cgi_get, getcache
from cantools.db import get_model
from cantools import config

def response():
	action = cgi_get("action", choices=["db", "memcache", "pubsub"])
	if cgi_get("pw") != config.admin:
		fail("wrong");
	if action == "db":
		mod = get_model(cgi_get("modelName"))
		limit = cgi_get("limit")
		offset = cgi_get("offset")
		succeed([d.data() for d in mod.query().fetch(limit, offset)])
	elif action == "memcache":
		succeed(getcache())
	elif action == "pubsub":
		succeed({ "host": config.pubsub.host, "port": config.pubsub.port })

respond(response)<commit_msg>/admin handler: schema response for db action requests that don't specify a model<commit_after>from cantools.web import respond, success, fail, cgi_get, getcache
from cantools.db import get_model, get_schema
from cantools import config

def response():
	action = cgi_get("action", choices=["db", "memcache", "pubsub"])
	if cgi_get("pw") != config.admin:
		fail("wrong");
	if action == "db":
		mname = cgi_get("modelName", required=False)
		if mname:
			mod = get_model(cgi_get("modelName"))
			limit = cgi_get("limit")
			offset = cgi_get("offset")
			succeed([d.data() for d in mod.query().fetch(limit, offset)])
		succeed(get_schema())
	elif action == "memcache":
		succeed(getcache())
	elif action == "pubsub":
		succeed({ "host": config.pubsub.host, "port": config.pubsub.port })

respond(response)<|endoftext|>


<commit_before><commit_msg>Workaround for a different sbt output<commit_after><|endoftext|>


<commit_before>import sys
sys.path.append('/home/jwalker/dynamics/python/atmos-tools')
sys.path.append('/home/jwalker/dynamics/python/atmos-read')

import xray
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
import collections
import pandas as pd
import scipy.signal
import atmos as atm
import precipdat
import merra
import indices

# ----------------------------------------------------------------------
isave = False
exts = ['png', 'eps']
index = collections.OrderedDict()

datadir = atm.homedir() + 'datastore/merra/daily/'
vimtfile = datadir + 'merra_vimt_ps-300mb_apr-sep_1979-2014.nc'
precipfile = datadir + 'merra_precip_40E-120E_60S-60N_days91-274_1979-2014.nc'
cmapfile = atm.homedir() + 'datastore/cmap/cmap.precip.pentad.mean.nc'
eraIfile = atm.homedir() + ('datastore/era_interim/analysis/'
                            'era_interim_JJAS_60-100E_index.csv')
ocifile = datadir + 'merra_u850_40E-120E_60S-60N_1979-2014.nc'

# Lat-lon box for WLH method
lon1, lon2 = 60, 100
lat1, lat2 = 10, 30

# ----------------------------------------------------------------------
# HOWI index (Webster and Fasullo 2003)
maxbreak = 10
with xray.open_dataset(vimtfile) as ds:
    ds.load()
for npts in [50, 100]:
    howi, _ = indices.onset_HOWI(ds['uq_int'], ds['vq_int'], npts,
                                 maxbreak=maxbreak)
    howi.attrs['title'] = 'HOWI (N=%d)' % npts
    index['HOWI_%d' % npts] = howi

# ----------------------------------------------------------------------
# Wang & LinHo method

def get_onset_WLH(years, days, pcp_sm, threshold, titlestr, pentad=True,
                  pcp_jan=None):
    nyears = len(years)
    i_onset = np.zeros(nyears)
    i_retreat = np.zeros(nyears)
    i_peak = np.zeros(nyears)
    for y, year in enumerate(years):
        vals = indices.onset_WLH_1D(pcp_sm[y], threshold, precip_jan=pcp_jan)
        i_onset[y], i_retreat[y], i_peak[y] = vals

    # Convert from pentads to day of year
    if pentad:
        d_onset = [int(atm.pentad_to_jday(i, pmin=0)) for i in i_onset]
        d_retreat = [int(atm.pentad_to_jday(i, pmin=0)) for i in i_retreat]
    else:
        d_onset = [np.nan if np.isnan(i) else days[int(i)] for i in i_onset]
        d_retreat = [np.nan if np.isnan(i) else days[int(i)] for i in i_retreat]

    # Pack into Dataset
    index = xray.Dataset()
    days = xray.DataArray(days, {'day' : days})
    years = xray.DataArray(years, {'year' : years})
    index['tseries'] = xray.DataArray(pcp_sm, dims=['year', 'day'],
                                      coords={'year' : years, 'day': days})
    index['onset'] = xray.DataArray(d_onset, coords={'year' : years})
    index['retreat'] = xray.DataArray(d_retreat, coords={'year' : years})
    index.attrs['title'] = titlestr
    return index

# Threshold and smoothing parameters
threshold = 5.0
kmax = 12
nroll = {'CMAP' : 3, 'MERRA_MFC' : 7, 'MERRA_PRECIP' : 7}

# Read CMAP pentad precip
cmap = precipdat.read_cmap(cmapfile)
cmapbar = atm.mean_over_geobox(cmap, lat1, lat2, lon1, lon2)
cmapdays = [atm.pentad_to_jday(p, pmin=1) for p in cmap.pentad.values]

# MERRA moisture flux convergence
print('Calculating MFC')
mfc = atm.moisture_flux_conv(ds['uq_int'], ds['vq_int'], already_int=True)
mfcbar = atm.mean_over_geobox(mfc, lat1, lat2, lon1, lon2)

# MERRA precip
print('Reading MERRA precip ' + precipfile)
with xray.open_dataset(precipfile) as dsprecip:
    precip = dsprecip['PRECTOT']
    precipbar = atm.mean_over_geobox(precip, lat1, lat2, lon1, lon2)

# Compute indices for each dataset
for name in ['CMAP', 'MERRA_MFC', 'MERRA_PRECIP']:
    print('****' + name + '******')
    if name == 'CMAP':
        pcp = cmapbar
        days = cmapdays
        pentad = True
        precip_jan = None # Calculate from pentad data
    elif name == 'MERRA_MFC':
        pcp = mfcbar
        days = mfcbar.day.values
        pentad = False
        precip_jan = 0.0 # Use zero for now
    elif name == 'MERRA_PRECIP':
        pcp = precipbar
        days = precipbar.day.values
        pentad = False
        precip_jan = 0.0 # Use zero for now

    years = pcp.year.values

    key = 'WLH_%s_kmax%d' % (name, kmax)
    print(key)
    pcp_sm, Rsq = atm.fourier_smooth(pcp, kmax)
    index[key] = get_onset_WLH(years, days, pcp_sm, threshold, key, pentad,
                               precip_jan)

    # Smooth with rolling mean
    key = 'WLH_%s_nroll%d' % (name, nroll[name])
    print(key)
    pcp_sm = np.zeros(pcp.shape)
    for y in range(pcp.shape[0]):
        pcp_sm[y] = pd.rolling_mean(pcp[y].values, nroll[name],
                                    center=True)
    index[key] = get_onset_WLH(years, days, pcp_sm, threshold, key, pentad,
                               precip_jan)

    # Unsmoothed pentad timeserires
    key = 'WLH_%s_unsmth' % name
    print(key)
    index[key] = get_onset_WLH(years, days, pcp, threshold, key, pentad,
                               precip_jan)

# ----------------------------------------------------------------------
# OCI index (Wang et al 2009)
with xray.open_dataset(ocifile) as ds:
    u850 = ds['U'].load()
u850 = u850.rename({'Year' : 'year', 'Day' : 'day'})
index['OCI'] = indices.onset_OCI(u850, yearnm='year', daynm='day')

# ----------------------------------------------------------------------
# Summary plots
for key in index.keys():
    ind = index[key]
    indices.summarize_indices(ind.year, ind.onset, ind.retreat, ind.title)
    if isave:
        for ext in exts:
            atm.savefigs('summary_' + key + '_', ext)
        plt.close('all')

# Plot daily timeseries for each year
keys = index.keys()
for key in keys:
    indices.plot_index_years(index[key], suptitle=key)
    if isave:
        for ext in exts:
            atm.savefigs('tseries_' + key + '_', ext)
        plt.close('all')

# Compare onset indices to each other
short = { 'HOWI_50' : 'HOWI_50',
          'HOWI_100' : 'HOWI_100',
          'WLH_CMAP_kmax12' : 'W_C_k12',
          'WLH_CMAP_nroll3' : 'W_C_n3',
          'WLH_CMAP_unsmth' : 'W_C_u',
          'WLH_MERRA_MFC_kmax12' : 'W_MM_k12',
          'WLH_MERRA_MFC_nroll7' : 'W_MM_n7',
          'WLH_MERRA_MFC_unsmth' : 'W_MM_u',
          'WLH_MERRA_PRECIP_kmax12' : 'W_MP_k12',
          'WLH_MERRA_PRECIP_nroll7' : 'W_MP_n7',
          'WLH_MERRA_PRECIP_unsmth' : 'W_MP_u',
          'OCI' : 'OCI'}

# Subset of keys to include in correlation calcs
keys = ['HOWI_100', 'HOWI_50', 'OCI', 'WLH_CMAP_kmax12', 'WLH_CMAP_nroll3',
        'WLH_MERRA_PRECIP_nroll7']
shortkeys = [short[key] for key in keys]
years = index[keys[0]].year.values
onset = np.reshape(index[keys[0]].onset.values, (len(years), 1))
for key in keys[1:]:
    ind = np.reshape(index[key].onset.values, (len(years), 1))
    onset = np.concatenate([onset, ind], axis=1)
onset = pd.DataFrame(onset, index=years, columns=shortkeys)

# Plot matrix of scatter plots with correlation coeffs
titlestr = 'Onset Day - Scatter Plots and Correlation Coefficients'
atm.scatter_matrix(onset, corr_fmt='.2f', corr_pos=(0.1, 0.85), figsize=(16,10),
                   suptitle=titlestr)

if isave:
    for ext in exts:
        atm.savefigs('onset_scatter', ext)
        plt.close('all')

# ----------------------------------------------------------------------
# Monsoon strength indices

def detrend(vals, index):
    vals_det = scipy.signal.detrend(vals)
    vals_det = vals_det / np.std(vals_det)
    output = pd.Series(vals_det, index=index)
    return output

mfc_JJAS = atm.subset(mfcbar, 'day', atm.season_days('JJAS'))
mfc_JJAS = mfc_JJAS.mean(dim='day')
mfc_det = detrend(mfc_JJAS.values, mfc_JJAS.year)

era = pd.read_csv(eraIfile, index_col=0)
era_det = detrend(era.values.flatten(), era.index)

strength = mfc_JJAS.to_series().to_frame(name='MERRA')
strength['MERRA_DET'] = mfc_det
strength['ERAI'] = era
strength['ERAI_DET'] = era_det

strength.plot()
plt.title('JJAS Monsoon Strength')
plt.xlabel('Year')
plt.ylabel('Index (mm/day)')

atm.scatter_matrix(strength)

# ----------------------------------------------------------------------
# Plotting timeseries together

data = xray.Dataset()
keys = index.keys()
for k in [1, 6, 9, 11]:
    key = keys[k]
    data[short[key]] = index[key]['tseries']

key_onset = 'HOWI_100'
d_onset = index[key_onset]['onset'].values
suptitle = key_onset + ' Onset'
indices.plot_tseries_together(data, onset=d_onset, suptitle=suptitle)
<commit_msg>Correlation coefficients between daily timeseries of two indices<commit_after>import sys
sys.path.append('/home/jwalker/dynamics/python/atmos-tools')
sys.path.append('/home/jwalker/dynamics/python/atmos-read')

import xray
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
import collections
import pandas as pd
import scipy.signal
import atmos as atm
import precipdat
import merra
import indices

# ----------------------------------------------------------------------
isave = False
exts = ['png', 'eps']
index = collections.OrderedDict()

datadir = atm.homedir() + 'datastore/merra/daily/'
vimtfile = datadir + 'merra_vimt_ps-300mb_apr-sep_1979-2014.nc'
precipfile = datadir + 'merra_precip_40E-120E_60S-60N_days91-274_1979-2014.nc'
cmapfile = atm.homedir() + 'datastore/cmap/cmap.precip.pentad.mean.nc'
eraIfile = atm.homedir() + ('datastore/era_interim/analysis/'
                            'era_interim_JJAS_60-100E_index.csv')
ocifile = datadir + 'merra_u850_40E-120E_60S-60N_1979-2014.nc'
ttfile = datadir + 'merra_T200-600_apr-sep_1979-2014.nc'

# Lat-lon box for WLH method
lon1, lon2 = 60, 100
lat1, lat2 = 10, 30

# ----------------------------------------------------------------------
# HOWI index (Webster and Fasullo 2003)
maxbreak = 10
with xray.open_dataset(vimtfile) as ds:
    ds.load()
for npts in [50, 100]:
    howi, _ = indices.onset_HOWI(ds['uq_int'], ds['vq_int'], npts,
                                 maxbreak=maxbreak)
    howi.attrs['title'] = 'HOWI (N=%d)' % npts
    index['HOWI_%d' % npts] = howi

# ----------------------------------------------------------------------
# Wang & LinHo method

def get_onset_WLH(years, days, pcp_sm, threshold, titlestr, pentad=True,
                  pcp_jan=None):
    nyears = len(years)
    i_onset = np.zeros(nyears)
    i_retreat = np.zeros(nyears)
    i_peak = np.zeros(nyears)
    for y, year in enumerate(years):
        vals = indices.onset_WLH_1D(pcp_sm[y], threshold, precip_jan=pcp_jan)
        i_onset[y], i_retreat[y], i_peak[y] = vals

    # Convert from pentads to day of year
    if pentad:
        d_onset = [int(atm.pentad_to_jday(i, pmin=0)) for i in i_onset]
        d_retreat = [int(atm.pentad_to_jday(i, pmin=0)) for i in i_retreat]
    else:
        d_onset = [np.nan if np.isnan(i) else days[int(i)] for i in i_onset]
        d_retreat = [np.nan if np.isnan(i) else days[int(i)] for i in i_retreat]

    # Pack into Dataset
    index = xray.Dataset()
    days = xray.DataArray(days, {'day' : days})
    years = xray.DataArray(years, {'year' : years})
    index['tseries'] = xray.DataArray(pcp_sm, dims=['year', 'day'],
                                      coords={'year' : years, 'day': days})
    index['onset'] = xray.DataArray(d_onset, coords={'year' : years})
    index['retreat'] = xray.DataArray(d_retreat, coords={'year' : years})
    index.attrs['title'] = titlestr
    return index

# Threshold and smoothing parameters
threshold = 5.0
kmax = 12
nroll = {'CMAP' : 3, 'MERRA_MFC' : 7, 'MERRA_PRECIP' : 7}

# Read CMAP pentad precip
cmap = precipdat.read_cmap(cmapfile)
cmapbar = atm.mean_over_geobox(cmap, lat1, lat2, lon1, lon2)
cmapdays = [atm.pentad_to_jday(p, pmin=1) for p in cmap.pentad.values]

# MERRA moisture flux convergence
print('Calculating MFC')
mfc = atm.moisture_flux_conv(ds['uq_int'], ds['vq_int'], already_int=True)
mfcbar = atm.mean_over_geobox(mfc, lat1, lat2, lon1, lon2)

# MERRA precip
print('Reading MERRA precip ' + precipfile)
with xray.open_dataset(precipfile) as dsprecip:
    precip = dsprecip['PRECTOT']
    precipbar = atm.mean_over_geobox(precip, lat1, lat2, lon1, lon2)

# Compute indices for each dataset
for name in ['CMAP', 'MERRA_MFC', 'MERRA_PRECIP']:
    print('****' + name + '******')
    if name == 'CMAP':
        pcp = cmapbar
        days = cmapdays
        pentad = True
        precip_jan = None # Calculate from pentad data
    elif name == 'MERRA_MFC':
        pcp = mfcbar
        days = mfcbar.day.values
        pentad = False
        precip_jan = 0.0 # Use zero for now
    elif name == 'MERRA_PRECIP':
        pcp = precipbar
        days = precipbar.day.values
        pentad = False
        precip_jan = 0.0 # Use zero for now

    years = pcp.year.values

    key = 'WLH_%s_kmax%d' % (name, kmax)
    print(key)
    pcp_sm, Rsq = atm.fourier_smooth(pcp, kmax)
    index[key] = get_onset_WLH(years, days, pcp_sm, threshold, key, pentad,
                               precip_jan)

    # Smooth with rolling mean
    key = 'WLH_%s_nroll%d' % (name, nroll[name])
    print(key)
    pcp_sm = np.zeros(pcp.shape)
    for y in range(pcp.shape[0]):
        pcp_sm[y] = pd.rolling_mean(pcp[y].values, nroll[name],
                                    center=True)
    index[key] = get_onset_WLH(years, days, pcp_sm, threshold, key, pentad,
                               precip_jan)

    # Unsmoothed pentad timeserires
    key = 'WLH_%s_unsmth' % name
    print(key)
    index[key] = get_onset_WLH(years, days, pcp, threshold, key, pentad,
                               precip_jan)

# ----------------------------------------------------------------------
# OCI index (Wang et al 2009)
with xray.open_dataset(ocifile) as ds:
    u850 = ds['U'].load()
u850 = u850.rename({'Year' : 'year', 'Day' : 'day'})
index['OCI'] = indices.onset_OCI(u850, yearnm='year', daynm='day')
index['OCI'].attrs['title'] = 'OCI'

# ----------------------------------------------------------------------
# TT index (Goswami et al 2006)
with xray.open_dataset(ttfile) as ds:
    T = ds['T'].load()
north=(5, 30, 40, 100)
south=(-15, 5, 40, 100)
index['TT'] = indices.onset_TT(T, north=north, south=south)
index['TT'].attrs['title'] = 'TT'

# ----------------------------------------------------------------------
# Summary plots
for key in index.keys():
    ind = index[key]
    if 'retreat' in ind.keys():
        retreat = ind.retreat
    else:
        retreat = None
    indices.summarize_indices(ind.year, ind.onset, retreat, ind.title)
    if isave:
        for ext in exts:
            atm.savefigs('summary_' + key + '_', ext)
        plt.close('all')

# Plot daily timeseries for each year
keys = index.keys()
for key in keys:
    indices.plot_index_years(index[key], suptitle=key)
    if isave:
        for ext in exts:
            atm.savefigs('tseries_' + key + '_', ext)
        plt.close('all')

# Compare onset indices to each other
short = { 'HOWI_50' : 'HOWI_50',
          'HOWI_100' : 'HOWI_100',
          'WLH_CMAP_kmax12' : 'W_C_k12',
          'WLH_CMAP_nroll3' : 'W_C_n3',
          'WLH_CMAP_unsmth' : 'W_C_u',
          'WLH_MERRA_MFC_kmax12' : 'W_MM_k12',
          'WLH_MERRA_MFC_nroll7' : 'W_MM_n7',
          'WLH_MERRA_MFC_unsmth' : 'W_MM_u',
          'WLH_MERRA_PRECIP_kmax12' : 'W_MP_k12',
          'WLH_MERRA_PRECIP_nroll7' : 'W_MP_n7',
          'WLH_MERRA_PRECIP_unsmth' : 'W_MP_u',
          'OCI' : 'OCI',
          'TT' : 'TT'}

# Subset of keys to include in correlation calcs
keys = ['HOWI_100', 'HOWI_50', 'OCI', 'TT', 'WLH_CMAP_kmax12',
        'WLH_CMAP_nroll3', 'WLH_MERRA_PRECIP_nroll7']
shortkeys = [short[key] for key in keys]
years = index[keys[0]].year.values
onset = np.reshape(index[keys[0]].onset.values, (len(years), 1))
for key in keys[1:]:
    ind = np.reshape(index[key].onset.values, (len(years), 1))
    onset = np.concatenate([onset, ind], axis=1)
onset = pd.DataFrame(onset, index=years, columns=shortkeys)

# Plot matrix of scatter plots with correlation coeffs
titlestr = 'Onset Day - Scatter Plots and Correlation Coefficients'
atm.scatter_matrix(onset, corr_fmt='.2f', corr_pos=(0.1, 0.85), figsize=(16,10),
                   suptitle=titlestr)

if isave:
    for ext in exts:
        atm.savefigs('onset_scatter', ext)
        plt.close('all')

# ----------------------------------------------------------------------
# Monsoon strength indices

def detrend(vals, index):
    vals_det = scipy.signal.detrend(vals)
    vals_det = vals_det / np.std(vals_det)
    output = pd.Series(vals_det, index=index)
    return output

mfc_JJAS = atm.subset(mfcbar, 'day', atm.season_days('JJAS'))
mfc_JJAS = mfc_JJAS.mean(dim='day')
mfc_det = detrend(mfc_JJAS.values, mfc_JJAS.year)

era = pd.read_csv(eraIfile, index_col=0)
era_det = detrend(era.values.flatten(), era.index)

strength = mfc_JJAS.to_series().to_frame(name='MERRA')
strength['MERRA_DET'] = mfc_det
strength['ERAI'] = era
strength['ERAI_DET'] = era_det

strength.plot()
plt.title('JJAS Monsoon Strength')
plt.xlabel('Year')
plt.ylabel('Index (mm/day)')

atm.scatter_matrix(strength)

# ----------------------------------------------------------------------
# Plotting timeseries together

data = xray.Dataset()
keys = index.keys()
for k in [1, 6, 9, 11, 12]:
    key = keys[k]
    data[short[key]] = index[key]['tseries']

key_onset = 'HOWI_100'
d_onset = index[key_onset]['onset'].values
suptitle = key_onset + ' Onset'
indices.plot_tseries_together(data, onset=d_onset, suptitle=suptitle)

ind1 = data[keys[1]]
ind2 = data[keys[11]]

corr = np.zeros(years.shape)
for y, year in enumerate(years):
    df = atm.subset(ind1, yearnm, year).to_series().to_frame(name=ind1.name)
    df[ind2.name] = atm.subset(ind2, yearnm, year).to_series()
    corr[y] = df.corr().as_matrix()[0, 1]
plt.figure()
plt.bar(years, corr)
plt.xlabel('Year')
plt.ylabel('Correlation Coefficient')
plt.title(ind1.name + ' correlation with ' + ind2.name)
<|endoftext|>


<commit_before>#!/usr/bin/env python2.7
"""
Baseline comparison of calliing needs 1000 genomes sample graphs
which are cosntructed from 1000 genomes vcf data.  To build these,
we also need the entire chromosome fasta sequence for each region.
This script downloads these from ucsc into one file that can be
used by vg construct.

mostly cribbed from:
https://github.com/glennhickey/vg2sg/blob/master/data/fetch1kgpRegion.py

"""

import argparse, sys, os, os.path, random, subprocess, shutil, itertools, glob

def parse_args(args):
    parser = argparse.ArgumentParser(description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)

    # General options
    parser.add_argument("--out_fa", type=str, default="data/g1kvcf/chrom.fa",
                        help="output directory")
    parser.add_argument("--chroms", type=str, default="5,6,13,17,19",
                        help="comma-separated list of chroms to download")
    parser.add_argument("--leaveChr", action="store_true",
                        help="dont strip chr from sequence names")
    args = args[1:]

    return parser.parse_args(args)

def downloadChrom(chrom, options, assembly = "hg38"):
    """ concatenate whole chromosome onto output file
    """
    # This is the URL to get it from
    fa_url = ("http://hgdownload.cse.ucsc.edu/goldenPath/{}/"
              "chromosomes/chr{}.fa.gz").format(assembly, chrom)

    # download and uncompress
    cmd = "curl {} | zcat".format(fa_url)
    if not options.leaveChr:
        # take chr out of sequence names
        cmd += " | sed \"s/chr{}/{}/\" ".format(chrom, chrom)
    # convert to upper case
    cmd += " | sed \"s/a/A/g;s/c/C/g;s/g/G/g;s/t/T/g;s/n/N/g\" "
    # tack onto output
    cmd += " >> {}".format(options.out_fa)
    os.system(cmd)

def main(args):

    options = parse_args(args)

    # make sure output directory exists and file doesn't
    if not os.path.isdir(os.path.dirname(options.out_fa)):
        os.makedirs(os.path.dirname(options.out_fa))
    os.system("rm -f {}".format(options.out_fa))

    # append each chrom to file
    for chrom in options.chroms.split(","):
        downloadChrom(chrom, options)

    #make index
    os.system("samtools faidx {}".format(options.out_fa))

if __name__ == "__main__" :
    sys.exit(main(sys.argv))
<commit_msg>Don't upper-case FASTA headers<commit_after>#!/usr/bin/env python2.7
"""
Baseline comparison of calliing needs 1000 genomes sample graphs
which are cosntructed from 1000 genomes vcf data.  To build these,
we also need the entire chromosome fasta sequence for each region.
This script downloads these from ucsc into one file that can be
used by vg construct.

mostly cribbed from:
https://github.com/glennhickey/vg2sg/blob/master/data/fetch1kgpRegion.py

"""

import argparse, sys, os, os.path, random, subprocess, shutil, itertools, glob

def parse_args(args):
    parser = argparse.ArgumentParser(description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)

    # General options
    parser.add_argument("--out_fa", type=str, default="data/g1kvcf/chrom.fa",
                        help="output directory")
    parser.add_argument("--chroms", type=str, default="5,6,13,17,19",
                        help="comma-separated list of chroms to download")
    parser.add_argument("--leaveChr", action="store_true",
                        help="dont strip chr from sequence names")
    args = args[1:]

    return parser.parse_args(args)

def downloadChrom(chrom, options, assembly = "hg38"):
    """ concatenate whole chromosome onto output file
    """
    # This is the URL to get it from
    fa_url = ("http://hgdownload.cse.ucsc.edu/goldenPath/{}/"
              "chromosomes/chr{}.fa.gz").format(assembly, chrom)

    # download and uncompress
    cmd = "curl {} | zcat".format(fa_url)
    if not options.leaveChr:
        # take chr out of sequence names
        cmd += " | sed \"s/chr{}/{}/\" ".format(chrom, chrom)
    # convert non-headers to upper case
    cmd += " | sed \"/^>/!s/a/A/g;/^>/!s/c/C/g;/^>/!s/g/G/g;/^>/!s/t/T/g;/^>/!s/n/N/g\" "
    # tack onto output
    cmd += " >> {}".format(options.out_fa)
    os.system(cmd)

def main(args):

    options = parse_args(args)

    # make sure output directory exists and file doesn't
    if not os.path.isdir(os.path.dirname(options.out_fa)):
        os.makedirs(os.path.dirname(options.out_fa))
    os.system("rm -f {}".format(options.out_fa))

    # append each chrom to file
    for chrom in options.chroms.split(","):
        downloadChrom(chrom, options)

    #make index
    os.system("samtools faidx {}".format(options.out_fa))

if __name__ == "__main__" :
    sys.exit(main(sys.argv))
<|endoftext|>


<commit_before>import yaml
import os
from pprint import pprint

def _merge_yaml(ret, data):
    '''
    Merge two yaml dicts together
    '''
    if 'sysctl' not in ret:
        ret['sysctl'] = {}
    for topkey in data.get('sysctl', {}):
        if topkey not in ret['sysctl']:
            ret['sysctl'][topkey] = {}
        ret['sysctl'][topkey].update(data['sysctl'][topkey])
    return ret

def _get_yaml(dirname):
    ret = {}
    try:
        for yamlpath in os.listdir(dirname):
            if yamlpath.endswith('.yaml') or yamlpath.endswith('.yml'):
                with open(os.path.join(dirname, yamlpath)) as fh_:
                    data = yaml.safe_load(fh_)
                _merge_yaml(ret, data)
    except:
        return {}
    return ret

if __name__ == "__main__":
    pprint(_get_yaml('.'))<commit_msg>update<commit_after><|endoftext|>


<commit_before>from __future__ import absolute_import
from celery.signals import setup_logging
import orchestrator.logger

__version__ = '0.3.10'
__author__ = 'sukrit'

orchestrator.logger.init_logging()
setup_logging.connect(orchestrator.logger.init_celery_logging)
<commit_msg>Prepare for next development version<commit_after>from __future__ import absolute_import
from celery.signals import setup_logging
import orchestrator.logger

__version__ = '0.4.0'
__author__ = 'sukrit'

orchestrator.logger.init_logging()
setup_logging.connect(orchestrator.logger.init_celery_logging)
<|endoftext|>