<commit_before>from sts.input_traces.fingerprints import *
from sts.replay_event import *
import sts.control_flow
import logging
import time
import math
from sts.util.convenience import find_index
log = logging.getLogger("event_dag")

class EventDag(object):
  # We peek ahead this many seconds after the timestamp of the subseqeunt
  # event
  # TODO(cs): be smarter about this -- peek() too far, and peek()'ing not far
  # enough can both have negative consequences
  _peek_seconds = 3.0
  # If we prune a failure, make sure that the subsequent
  # recovery doesn't occur
  _failure_types = set([SwitchFailure, LinkFailure, ControllerFailure, ControlChannelBlock])
  # NOTE: we treat failure/recovery as an atomic pair, since it doesn't make
  # much sense to prune a recovery event
  _recovery_types = set([SwitchRecovery, LinkRecovery, ControllerRecovery, ControlChannelUnblock])
  # For now, we're ignoring these input types, since their dependencies with
  # other inputs are too complicated to model
  # TODO(cs): model these!
  _ignored_input_types = set([DataplaneDrop, DataplanePermit, HostMigration])

  '''A collection of Event objects. EventDags are primarily used to present a
  view of the underlying events with some subset of the input events pruned
  '''
  def __init__(self, events, is_view=False, prefix_trie=None,
               label2event=None, ignore_unsupported_input_types=False,
               mark_invalid_input_sequences=False):
    '''events is a list of EventWatcher objects. Refer to log_parser.parse to
    see how this is assembled.'''
    if ignore_unsupported_input_types:
      self._events_list = [ e for e in events
                            if events not in self._ignored_input_types ]
    else:
      self._events_list = events
    self._populate_indices(label2event)

    # Fill in domain knowledge about valid input
    # sequences (e.g. don't prune failure without pruning recovery.)
    # Only do so if this isn't a view of a previously computed DAG
    # TODO(cs): there is probably a cleaner way to implement views
    if not is_view and mark_invalid_input_sequences:
      try:
        import pytrie
      except ImportError:
        raise RuntimeError("Need to install pytrie: `sudo pip install pytrie`")
      self._mark_invalid_input_sequences()
      prefix_trie = pytrie.Trie()
    # The prefix trie stores lists of input events as keys,
    # and lists of both input and internal events as values
    # Note that we pass the trie around between DAG views
    self._prefix_trie = prefix_trie

  def _populate_indices(self, label2event):
    self._event_to_index = {
      e : i
      for i, e in enumerate(self._events_list)
    }
    # Optimization: only compute label2event once (at the first unpruned
    # initialization)
    # TODO(cs): need to ensure that newly added events get labeled
    # uniquely. (Easy, but inelegant way: set the label generator's
    # initial value to max(labels) + 1)
    if label2event is None:
      self._label2event = {
        event.label : event
        for event in self._events_list
      }
    else:
      self._label2event = label2event

  @property
  def events(self):
    '''Return the events in the DAG'''
    return self._events_list

  @property
  def event_watchers(self):
    '''Return a generator of the EventWatchers in the DAG'''
    return map(EventWatcher, self._events_list)

  def _remove_event(self, event):
    ''' Recursively remove the event and its dependents '''
    if event in self._event_to_index:
      list_idx = self._event_to_index[event]
      del self._event_to_index[event]
      self._event_list.pop(list_idx)

    # Note that dependent_labels only contains dependencies between input
    # events. We run peek() to infer dependencies with internal events
    for label in event.dependent_labels:
      if label in self._label2event:
        dependent_event = self._label2event[label]
        if dependent_event in self._event_to_index:
          self._remove_event(dependent_event)

  def remove_events(self, ignored_portion, simulation):
    ''' Mutate the DAG: remove all input events in ignored_inputs,
    as well all of their dependent input events'''
    # Note that we treat failure/recovery as an atomic pair, so we don't prune
    # recovery events on their own
    for event in [ e for e in ignored_portion
                   if (isinstance(e, InputEvent) and
                       type(e) not in self._recovery_types) ]:
      self._remove_event(event)
    # Now run peek() to hide the internal events that will no longer occur
    # Note that causal dependencies change depending on what the prefix is!
    # So we have to run peek() once per prefix
    self.peek(simulation)

  def ignore_portion(self, ignored_portion, simulation):
    ''' Return a view of the dag with ignored_portion and its dependents
    removed'''
    dag = EventDag(list(self._events_list), is_view=True,
                   prefix_trie=self._prefix_trie,
                   label2event=self._label2event)
    # TODO(cs): potentially some redundant computation here
    dag.remove_events(ignored_portion, simulation)
    return dag

  def split_inputs(self, split_ways):
    ''' Split our events into split_ways separate lists '''
    events = self._events_list
    if len(events) == 0:
      return [[]]
    if split_ways == 1:
      return [events]
    if split_ways < 1 or split_ways > len(events):
      raise ValueError("Invalid split ways %d" % split_ways)

    splits = []
    split_interval = int(math.ceil(len(events) * 1.0 / split_ways))
    start_idx = 0
    split_idx = start_idx + split_interval
    while start_idx < len(events):
      splits.append(events[start_idx:split_idx])
      start_idx = split_idx
      # Account for odd numbered splits -- if we're about to eat up
      # all elements even though we will only have added split_ways-1
      # splits, back up the split interval by 1
      if (split_idx + split_interval >= len(events) and
          len(splits) == split_ways - 2):
        split_interval -= 1
      split_idx += split_interval
    return splits

  def peek(self, simulation):
    ''' Infer which internal events are/aren't going to occur, '''
    # TODO(cs): optimization: write the prefix trie to a file, in case we want to run
    # FindMCS again?
    input_events = [ e for e in self._events_list if isinstance(e, InputEvent) ]
    if len(input_events) == 0:
      # Postcondition: input_events[-1] is not None
      #                and self._events_list[-1] is not None
      return

    # Note that we recompute wait times for every view, since the set of
    # inputs and intervening expected internal events changes
    def get_wait_times(input_events):
      event2wait_time = {}
      for i in xrange(0, len(input_events)-1):
        current_input = input_events[i]
        next_input = input_events[i+1]
        wait_time = next_input.time.as_float() + self._peek_seconds
        event2wait_time[current_input] = wait_time
      # For the last event, we wait until the last internal event
      last_wait_time = self._events_list[-1].time.as_float() + self._peek_seconds
      event2wait_time[input_events[-1]] = last_wait_time
      return event2wait_time

    log.debug("Computing wait times")
    event2wait_time = get_wait_times(input_events)

    # Also compute the internal events that we expect for each interval between
    # input events
    def get_expected_internal_events(input_events):
      input_to_exected_events = {}
      for i in xrange(0, len(input_events)-1):
        # Infer the internal events that we expect
        current_input = input_events[i]
        current_input_idx = self._event_to_index[current_input]
        next_input = input_events[i+1]
        next_input_idx = self._event_to_index[next_input]
        expected_internal_events = \
                self._events_list[current_input_idx+1:next_input_idx]
        input_to_expected_events[current_input] = expected_internal_events
      # The last input's expected internal events are anything that follow it
      # in the log.
      last_input = input_events[-1]
      last_input_idx = self._event_to_index[last_input]
      input_to_expected_events[last_input] = self._events_list[last_input_idx:]
      return input_to_exected_events

    log.debug("Computing expected internal events")
    input_to_expected_events = get_expected_internal_events(input_events)

    # Now, play the execution forward iteratively for each input event, and
    # record what internal events happen between the injection and the
    # wait_time

    # Initilize current_input_prefix to the longest_match prefix we've
    # inferred previously (or [] if this is an entirely new prefix)
    current_input_prefix = self._prefix_trie\
                               .longest_prefix(input_events, default=[])

    # The value is both internal events and input events (values of the trie)
    # leading up to, but not including the next input following the tail of the prefix
    inferred_events = self._prefix_trie\
                          .longest_prefix_value(input_events, default=[])

    # current_input is the next input after the tail of the prefix
    if current_input_prefix == []:
      current_input_idx = 0
    else:
      current_input_idx = input_events.index(current_input_prefix[-1]) + 1

    while current_input_idx < len(input_events):
      log.debug("Pruning input index %d" % current_input_idx)
      current_input = input_events[current_input_idx]
      expected_internal_events = input_to_expected_events[current_input]
      # Optimization: if no internal events occured between this input and the
      # next, no need to peek()
      if expected_internal_events == []:
        log.debug("Optimization: no expected internal events")
        newly_inferred_events = [current_input]
      else:
        # Now actually do the peek()'ing! First replay the prefix
        # plus the next input
        prefix_dag = EventDag(inferred_events + [current_input])
        replayer = control_flow.Replayer(prefix_dag)
        log.debug("Replaying prefix")
        replayer.simulate(simulation)

        # Directly after the last input has been injected, flush the internal
        # event buffers in case there were unaccounted internal events
        # Note that there isn't a race condition between flush()'ing and
        # incoming internal events, since sts is single-threaded
        simulation.god_scheduler.flush()
        simulation.controller_sync_callback.flush()

        # Now set all internal event buffers (GodScheduler for
        # ControlMessageReceives and ReplaySyncCallback for state changes)
        # to "pass through + record" by defining event handlers
        newly_inferred_events = []
        def receipt_pass_through(receipt_event):
          pending_receipt = receipt_event.pending_receipt
          # Pass through
          simulation.god_scheduler.schedule(pending_receipt)
          # Record
          replay_event = ControlMessageReceive(pending_receipt.dpid,
                                               pending_receipt.controller_id,
                                               pending_receipt.fingerprint.to_dict())
          newly_inferred_events.append(replay_event)

        simulation.god_scheduler.addListener(MessageReceipt,
                                             receipt_pass_through)

        def state_change_pass_through(state_change_event):
          state_change = state_change_event.pending_state_change
          # Pass through
          simulation.controller_sync_callback.gc_pending_state_change(state_change)
          # Record
          replay_event = ControllerStateChange(state_change.controller_id,
                                               state_change.time,
                                               state_change.fingerprint,
                                               state_change.name,
                                               state_change.value)
          newly_inferred_events.append(replay_event)

        simulation.controller_sync_callback.addListener(control_flow.StateChange,
                                                        state_change_pass_through)

        # Now sit tight for wait_seconds
        wait_seconds = event2wait_time[current_input]
        # Note that this is the monkey patched version of time.sleep
        log.debug("peek()'ing for %f seconds" % wait_seconds)
        time.sleep(wait_seconds)

        # Now turn off those listeners
        simulation.god_scheduler.removeListener(receipt_pass_through)
        simulation.controller_sync_callback.removeListener(state_change_pass_through)

        # Finally, insert current_input into the appropriate place in
        # inferred_events (and ignore internal events that come afterward)
        def match_fingerprints(newly_inferred_events):
          # Find the last internal event in expected_internal_events that
          # matches an event in newly_inferred_events. That is the new causal
          # parent of current_input
          expected_internal_events.reverse()
          inferred_fingerprints = set([e.fingerprint
                                       for e in newly_inferred_events])
          for expected in expected_internal_events:
            if expected.fingerprint in inferred_fingerprints:
              # We've found our insertion point.
              # Insert the input after the expected internal event, and all
              # internal events that come after it
              parent_index = find_index(lambda e: e.fingerprint == expected.fingerprint,
                                        newly_inferred_events)
              newly_inferred_events = newly_inferred_events[:parent_index+1]
              newly_inferred_events.append(current_input)
              return newly_inferred_events
          # Else, no expected internal event was observed, so just insert
          # current_input by itself
          return [current_input]

        log.debug("Matching fingerprints")
        log.debug("Expected: %s" % str(expected_internal_events))
        log.debug("Inferred: %s" % str(newly_inferred_events))
        newly_inferred_events = match_fingerprints(newly_inferred_events)

      # Update the trie for this prefix
      current_input_prefix.append(current_input)
      # Note that newly_inferred_events already includes current_input
      inferred_events += newly_inferred_events
      self._prefix_trie[current_input_prefix] = inferred_events

      current_input_idx += 1

    # Now that the new execution has been inferred,
    # present a view of ourselves that only includes the updated
    # events
    self._events_list = inferred_events
    self._populate_indices(self._label2event)

  def _mark_invalid_input_sequences(self):
    # Note: we treat each failure/recovery pair atomically, since it doesn't
    # make much sense to prune recovery events. Also note that that we will
    # never see two failures (for a particular node) in a row without an
    # interleaving recovery event
    fingerprint2previousfailure = {}

    # NOTE: mutates self._events
    for event in self._events_list:
      if type(event) in self._failure_types:
        # Insert it into the previous failure hash
        fingerprint2previousfailure[event.fingerprint] = event
      elif type(event) in self._recovery_types:
        # Check if there were any failure predecessors
        if event.fingerprint in fingerprint2previousfailure:
          failure = fingerprint2previousfailure[event.fingerprint]
          failure.dependent_labels.append(event.label)
      elif type(event) in self._ignored_input_types:
        raise RuntimeError("No support for %s dependencies" %
                            type(event).__name__)

class EventWatcher(object):
  '''EventWatchers watch events. This class can be used to wrap either
  InternalEvents or ExternalEvents to perform pre and post functionality.'''

  def __init__(self, event):
    self.event = event

  def run(self, simulation):
    self._pre()

    while not self.event.proceed(simulation):
      time.sleep(0.05)
      log.debug(".")

    self._post()

  def _pre(self):
    log.debug("Executing %s" % str(self.event))

  def _post(self):
    log.debug("Finished Executing %s" % str(self.event))

<commit_msg>Minor type errors in event_dag.py<commit_after>from sts.input_traces.fingerprints import *
from sts.replay_event import *
import sts.control_flow
import logging
import time
import math
from sts.util.convenience import find_index
log = logging.getLogger("event_dag")

class EventDag(object):
  # We peek ahead this many seconds after the timestamp of the subseqeunt
  # event
  # TODO(cs): be smarter about this -- peek() too far, and peek()'ing not far
  # enough can both have negative consequences
  _peek_seconds = 3.0
  # If we prune a failure, make sure that the subsequent
  # recovery doesn't occur
  _failure_types = set([SwitchFailure, LinkFailure, ControllerFailure, ControlChannelBlock])
  # NOTE: we treat failure/recovery as an atomic pair, since it doesn't make
  # much sense to prune a recovery event
  _recovery_types = set([SwitchRecovery, LinkRecovery, ControllerRecovery, ControlChannelUnblock])
  # For now, we're ignoring these input types, since their dependencies with
  # other inputs are too complicated to model
  # TODO(cs): model these!
  _ignored_input_types = set([DataplaneDrop, DataplanePermit, HostMigration])

  '''A collection of Event objects. EventDags are primarily used to present a
  view of the underlying events with some subset of the input events pruned
  '''
  def __init__(self, events, is_view=False, prefix_trie=None,
               label2event=None, ignore_unsupported_input_types=False,
               mark_invalid_input_sequences=False):
    '''events is a list of EventWatcher objects. Refer to log_parser.parse to
    see how this is assembled.'''
    if ignore_unsupported_input_types:
      self._events_list = [ e for e in events
                            if e not in self._ignored_input_types ]
    else:
      self._events_list = events
    self._populate_indices(label2event)

    # Fill in domain knowledge about valid input
    # sequences (e.g. don't prune failure without pruning recovery.)
    # Only do so if this isn't a view of a previously computed DAG
    # TODO(cs): there is probably a cleaner way to implement views
    if not is_view and mark_invalid_input_sequences:
      try:
        import pytrie
      except ImportError:
        raise RuntimeError("Need to install pytrie: `sudo pip install pytrie`")
      self._mark_invalid_input_sequences()
      prefix_trie = pytrie.Trie()
    # The prefix trie stores lists of input events as keys,
    # and lists of both input and internal events as values
    # Note that we pass the trie around between DAG views
    self._prefix_trie = prefix_trie

  def _populate_indices(self, label2event):
    self._event_to_index = {
      e : i
      for i, e in enumerate(self._events_list)
    }
    # Optimization: only compute label2event once (at the first unpruned
    # initialization)
    # TODO(cs): need to ensure that newly added events get labeled
    # uniquely. (Easy, but inelegant way: set the label generator's
    # initial value to max(labels) + 1)
    if label2event is None:
      self._label2event = {
        event.label : event
        for event in self._events_list
      }
    else:
      self._label2event = label2event

  @property
  def events(self):
    '''Return the events in the DAG'''
    return self._events_list

  @property
  def event_watchers(self):
    '''Return a generator of the EventWatchers in the DAG'''
    return map(EventWatcher, self._events_list)

  def _remove_event(self, event):
    ''' Recursively remove the event and its dependents '''
    if event in self._event_to_index:
      list_idx = self._event_to_index[event]
      del self._event_to_index[event]
      self._event_list.pop(list_idx)

    # Note that dependent_labels only contains dependencies between input
    # events. We run peek() to infer dependencies with internal events
    for label in event.dependent_labels:
      if label in self._label2event:
        dependent_event = self._label2event[label]
        if dependent_event in self._event_to_index:
          self._remove_event(dependent_event)

  def remove_events(self, ignored_portion, simulation):
    ''' Mutate the DAG: remove all input events in ignored_inputs,
    as well all of their dependent input events'''
    # Note that we treat failure/recovery as an atomic pair, so we don't prune
    # recovery events on their own
    for event in [ e for e in ignored_portion
                   if (isinstance(e, InputEvent) and
                       type(e) not in self._recovery_types) ]:
      self._remove_event(event)
    # Now run peek() to hide the internal events that will no longer occur
    # Note that causal dependencies change depending on what the prefix is!
    # So we have to run peek() once per prefix
    self.peek(simulation)

  def ignore_portion(self, ignored_portion, simulation):
    ''' Return a view of the dag with ignored_portion and its dependents
    removed'''
    dag = EventDag(list(self._events_list), is_view=True,
                   prefix_trie=self._prefix_trie,
                   label2event=self._label2event)
    # TODO(cs): potentially some redundant computation here
    dag.remove_events(ignored_portion, simulation)
    return dag

  def split_inputs(self, split_ways):
    ''' Split our events into split_ways separate lists '''
    events = self._events_list
    if len(events) == 0:
      return [[]]
    if split_ways == 1:
      return [events]
    if split_ways < 1 or split_ways > len(events):
      raise ValueError("Invalid split ways %d" % split_ways)

    splits = []
    split_interval = int(math.ceil(len(events) * 1.0 / split_ways))
    start_idx = 0
    split_idx = start_idx + split_interval
    while start_idx < len(events):
      splits.append(events[start_idx:split_idx])
      start_idx = split_idx
      # Account for odd numbered splits -- if we're about to eat up
      # all elements even though we will only have added split_ways-1
      # splits, back up the split interval by 1
      if (split_idx + split_interval >= len(events) and
          len(splits) == split_ways - 2):
        split_interval -= 1
      split_idx += split_interval
    return splits

  def peek(self, simulation):
    ''' Infer which internal events are/aren't going to occur, '''
    # TODO(cs): optimization: write the prefix trie to a file, in case we want to run
    # FindMCS again?
    input_events = [ e for e in self._events_list if isinstance(e, InputEvent) ]
    if len(input_events) == 0:
      # Postcondition: input_events[-1] is not None
      #                and self._events_list[-1] is not None
      return

    # Note that we recompute wait times for every view, since the set of
    # inputs and intervening expected internal events changes
    def get_wait_times(input_events):
      event2wait_time = {}
      for i in xrange(0, len(input_events)-1):
        current_input = input_events[i]
        next_input = input_events[i+1]
        wait_time = next_input.time.as_float() + self._peek_seconds
        event2wait_time[current_input] = wait_time
      # For the last event, we wait until the last internal event
      last_wait_time = self._events_list[-1].time.as_float() + self._peek_seconds
      event2wait_time[input_events[-1]] = last_wait_time
      return event2wait_time

    log.debug("Computing wait times")
    event2wait_time = get_wait_times(input_events)

    # Also compute the internal events that we expect for each interval between
    # input events
    def get_expected_internal_events(input_events):
      input_to_exected_events = {}
      for i in xrange(0, len(input_events)-1):
        # Infer the internal events that we expect
        current_input = input_events[i]
        current_input_idx = self._event_to_index[current_input]
        next_input = input_events[i+1]
        next_input_idx = self._event_to_index[next_input]
        expected_internal_events = \
                self._events_list[current_input_idx+1:next_input_idx]
        input_to_expected_events[current_input] = expected_internal_events
      # The last input's expected internal events are anything that follow it
      # in the log.
      last_input = input_events[-1]
      last_input_idx = self._event_to_index[last_input]
      input_to_expected_events[last_input] = self._events_list[last_input_idx:]
      return input_to_exected_events

    log.debug("Computing expected internal events")
    input_to_expected_events = get_expected_internal_events(input_events)

    # Now, play the execution forward iteratively for each input event, and
    # record what internal events happen between the injection and the
    # wait_time

    # Initilize current_input_prefix to the longest_match prefix we've
    # inferred previously (or [] if this is an entirely new prefix)
    current_input_prefix = self._prefix_trie\
                               .longest_prefix(input_events, default=[])

    # The value is both internal events and input events (values of the trie)
    # leading up to, but not including the next input following the tail of the prefix
    inferred_events = self._prefix_trie\
                          .longest_prefix_value(input_events, default=[])

    # current_input is the next input after the tail of the prefix
    if current_input_prefix == []:
      current_input_idx = 0
    else:
      current_input_idx = input_events.index(current_input_prefix[-1]) + 1

    while current_input_idx < len(input_events):
      log.debug("Pruning input index %d" % current_input_idx)
      current_input = input_events[current_input_idx]
      expected_internal_events = input_to_expected_events[current_input]
      # Optimization: if no internal events occured between this input and the
      # next, no need to peek()
      if expected_internal_events == []:
        log.debug("Optimization: no expected internal events")
        newly_inferred_events = [current_input]
      else:
        # Now actually do the peek()'ing! First replay the prefix
        # plus the next input
        prefix_dag = EventDag(inferred_events + [current_input])
        replayer = control_flow.Replayer(prefix_dag)
        log.debug("Replaying prefix")
        replayer.simulate(simulation)

        # Directly after the last input has been injected, flush the internal
        # event buffers in case there were unaccounted internal events
        # Note that there isn't a race condition between flush()'ing and
        # incoming internal events, since sts is single-threaded
        simulation.god_scheduler.flush()
        simulation.controller_sync_callback.flush()

        # Now set all internal event buffers (GodScheduler for
        # ControlMessageReceives and ReplaySyncCallback for state changes)
        # to "pass through + record" by defining event handlers
        newly_inferred_events = []
        def receipt_pass_through(receipt_event):
          pending_receipt = receipt_event.pending_receipt
          # Pass through
          simulation.god_scheduler.schedule(pending_receipt)
          # Record
          replay_event = ControlMessageReceive(pending_receipt.dpid,
                                               pending_receipt.controller_id,
                                               pending_receipt.fingerprint.to_dict())
          newly_inferred_events.append(replay_event)

        simulation.god_scheduler.addListener(MessageReceipt,
                                             receipt_pass_through)

        def state_change_pass_through(state_change_event):
          state_change = state_change_event.pending_state_change
          # Pass through
          simulation.controller_sync_callback.gc_pending_state_change(state_change)
          # Record
          replay_event = ControllerStateChange(state_change.controller_id,
                                               state_change.time,
                                               state_change.fingerprint,
                                               state_change.name,
                                               state_change.value)
          newly_inferred_events.append(replay_event)

        simulation.controller_sync_callback.addListener(control_flow.StateChange,
                                                        state_change_pass_through)

        # Now sit tight for wait_seconds
        wait_seconds = event2wait_time[current_input]
        # Note that this is the monkey patched version of time.sleep
        log.debug("peek()'ing for %f seconds" % wait_seconds)
        time.sleep(wait_seconds)

        # Now turn off those listeners
        simulation.god_scheduler.removeListener(receipt_pass_through)
        simulation.controller_sync_callback.removeListener(state_change_pass_through)

        # Finally, insert current_input into the appropriate place in
        # inferred_events (and ignore internal events that come afterward)
        def match_fingerprints(newly_inferred_events):
          # Find the last internal event in expected_internal_events that
          # matches an event in newly_inferred_events. That is the new causal
          # parent of current_input
          expected_internal_events.reverse()
          inferred_fingerprints = set([e.fingerprint
                                       for e in newly_inferred_events])
          for expected in expected_internal_events:
            if expected.fingerprint in inferred_fingerprints:
              # We've found our insertion point.
              # Insert the input after the expected internal event, and all
              # internal events that come after it
              parent_index = find_index(lambda e: e.fingerprint == expected.fingerprint,
                                        newly_inferred_events)
              newly_inferred_events = newly_inferred_events[:parent_index+1]
              newly_inferred_events.append(current_input)
              return newly_inferred_events
          # Else, no expected internal event was observed, so just insert
          # current_input by itself
          return [current_input]

        log.debug("Matching fingerprints")
        log.debug("Expected: %s" % str(expected_internal_events))
        log.debug("Inferred: %s" % str(newly_inferred_events))
        newly_inferred_events = match_fingerprints(newly_inferred_events)

      # Update the trie for this prefix
      current_input_prefix.append(current_input)
      # Note that newly_inferred_events already includes current_input
      inferred_events += newly_inferred_events
      self._prefix_trie[current_input_prefix] = inferred_events

      current_input_idx += 1

    # Now that the new execution has been inferred,
    # present a view of ourselves that only includes the updated
    # events
    self._events_list = inferred_events
    self._populate_indices(self._label2event)

  def _mark_invalid_input_sequences(self):
    # Note: we treat each failure/recovery pair atomically, since it doesn't
    # make much sense to prune recovery events. Also note that that we will
    # never see two failures (for a particular node) in a row without an
    # interleaving recovery event
    fingerprint2previousfailure = {}

    # NOTE: mutates self._events
    for event in self._events_list:
      if type(event) in self._failure_types:
        # Insert it into the previous failure hash
        fingerprint2previousfailure[event.fingerprint] = event
      elif type(event) in self._recovery_types:
        # Check if there were any failure predecessors
        if event.fingerprint in fingerprint2previousfailure:
          failure = fingerprint2previousfailure[event.fingerprint]
          failure.dependent_labels.append(event.label)
      #elif type(event) in self._ignored_input_types:
      #  raise RuntimeError("No support for %s dependencies" %
      #                      type(event).__name__)

class EventWatcher(object):
  '''EventWatchers watch events. This class can be used to wrap either
  InternalEvents or ExternalEvents to perform pre and post functionality.'''

  def __init__(self, event):
    self.event = event

  def run(self, simulation):
    self._pre()

    while not self.event.proceed(simulation):
      time.sleep(0.05)
      log.debug(".")

    self._post()

  def _pre(self):
    log.debug("Executing %s" % str(self.event))

  def _post(self):
    log.debug("Finished Executing %s" % str(self.event))

<|endoftext|>


<commit_before># coding=utf-8

import re
import os
import pprint
from collections import OrderedDict

from bs4 import BeautifulSoup

TEMPLATE = """\
import re
from collections import OrderedDict
data = """

TEMPLATE_END = """\

for lang, grps in data.iteritems():
    for grp in grps.iterkeys():
        if data[lang][grp]["pattern"]:
            data[lang][grp]["pattern"] = re.compile(data[lang][grp]["pattern"])
"""

SZ_FIX_DATA = {
    "eng": {
        "PartialWordsAlways": {
            u"°x°": u"%",
            u"compiete": u"complete",
            u"Âs": u"'s",
            u"ÃÂs": u"'s",
        },
        "WholeWords": {
            u"I'11": u"I'll",
            u"Tun": u"Run",
            u"pan'": u"part",
            u"al'": u"at",
            u"a re": u"are",
            u"wail'": u"wait",
            u"he)'": u"hey",
            u"He)'": u"Hey",
            u"Yea h": u"Yeah",
            u"yea h": u"yeah",
        }
    }
}

if __name__ == "__main__":
    cur_dir = os.path.dirname(os.path.realpath(__file__))
    xml_dir = os.path.join(cur_dir, "xml")
    file_list = os.listdir(xml_dir)

    data = {}

    for fn in file_list:
        if fn.endswith("_OCRFixReplaceList.xml"):
            lang = fn.split("_")[0]
            soup = BeautifulSoup(open(os.path.join(xml_dir, fn)), "xml")

            fetch_data = (
                    # group, item_name, pattern
                    ("WholeLines", "Line", None),
                    ("WholeWords", "Word", lambda d: (ur"(?um)\b(?:" + u"|".join([re.escape(k) for k in d.keys()])
                                                      + ur')\b') if d else None),
                    ("PartialWordsAlways", "WordPart", None),
                    ("PartialLines", "LinePart", lambda d: (ur"(?um)(?:(?<=\s)|(?<=^)|(?<=\b))(?:" +
                                                            u"|".join([re.escape(k) for k in d.keys()]) +
                                                            ur")(?:(?=\s)|(?=$)|(?=\b))") if d else None),
                    ("BeginLines", "Beginning", lambda d: (ur"(?um)^(?:"+u"|".join([re.escape(k) for k in d.keys()])
                                                           + ur')') if d else None),
                    ("EndLines", "Ending", lambda d: (ur"(?um)(?:" + u"|".join([re.escape(k) for k in d.keys()]) +
                                                      ur")$") if d else None,),
            )

            data[lang] = dict((grp, {"data": OrderedDict(), "pattern": None}) for grp, item_name, pattern in fetch_data)

            for grp, item_name, pattern in fetch_data:
                for grp_data in soup.find_all(grp):
                    for line in grp_data.find_all(item_name):
                        data[lang][grp]["data"][line["from"]] = line["to"]

                # add our own dictionaries
                if lang in SZ_FIX_DATA and grp in SZ_FIX_DATA[lang]:
                    data[lang][grp]["data"].update(SZ_FIX_DATA[lang][grp])

                if pattern:
                    data[lang][grp]["pattern"] = pattern(data[lang][grp]["data"])

    f = open(os.path.join(cur_dir, "data.py"), "w+")
    f.write(TEMPLATE)
    f.write(pprint.pformat(data, width=1))
    f.write(TEMPLATE_END)
    f.close()
<commit_msg>submod: update eng OCR fix data<commit_after># coding=utf-8

import re
import os
import pprint
from collections import OrderedDict

from bs4 import BeautifulSoup

TEMPLATE = """\
import re
from collections import OrderedDict
data = """

TEMPLATE_END = """\

for lang, grps in data.iteritems():
    for grp in grps.iterkeys():
        if data[lang][grp]["pattern"]:
            data[lang][grp]["pattern"] = re.compile(data[lang][grp]["pattern"])
"""

SZ_FIX_DATA = {
    "eng": {
        "PartialWordsAlways": {
            u"°x°": u"%",
            u"compiete": u"complete",
            u"Âs": u"'s",
            u"ÃÂs": u"'s",
            u"a/ion": u"ation",
        },
        "WholeWords": {
            u"I'11": u"I'll",
            u"Tun": u"Run",
            u"pan'": u"part",
            u"al'": u"at",
            u"a re": u"are",
            u"wail'": u"wait",
            u"he)'": u"hey",
            u"He)'": u"Hey",
            u"Yea h": u"Yeah",
            u"yea h": u"yeah",
            u"h is": u"his",
        }
    }
}

if __name__ == "__main__":
    cur_dir = os.path.dirname(os.path.realpath(__file__))
    xml_dir = os.path.join(cur_dir, "xml")
    file_list = os.listdir(xml_dir)

    data = {}

    for fn in file_list:
        if fn.endswith("_OCRFixReplaceList.xml"):
            lang = fn.split("_")[0]
            soup = BeautifulSoup(open(os.path.join(xml_dir, fn)), "xml")

            fetch_data = (
                    # group, item_name, pattern
                    ("WholeLines", "Line", None),
                    ("WholeWords", "Word", lambda d: (ur"(?um)\b(?:" + u"|".join([re.escape(k) for k in d.keys()])
                                                      + ur')\b') if d else None),
                    ("PartialWordsAlways", "WordPart", None),
                    ("PartialLines", "LinePart", lambda d: (ur"(?um)(?:(?<=\s)|(?<=^)|(?<=\b))(?:" +
                                                            u"|".join([re.escape(k) for k in d.keys()]) +
                                                            ur")(?:(?=\s)|(?=$)|(?=\b))") if d else None),
                    ("BeginLines", "Beginning", lambda d: (ur"(?um)^(?:"+u"|".join([re.escape(k) for k in d.keys()])
                                                           + ur')') if d else None),
                    ("EndLines", "Ending", lambda d: (ur"(?um)(?:" + u"|".join([re.escape(k) for k in d.keys()]) +
                                                      ur")$") if d else None,),
            )

            data[lang] = dict((grp, {"data": OrderedDict(), "pattern": None}) for grp, item_name, pattern in fetch_data)

            for grp, item_name, pattern in fetch_data:
                for grp_data in soup.find_all(grp):
                    for line in grp_data.find_all(item_name):
                        data[lang][grp]["data"][line["from"]] = line["to"]

                # add our own dictionaries
                if lang in SZ_FIX_DATA and grp in SZ_FIX_DATA[lang]:
                    data[lang][grp]["data"].update(SZ_FIX_DATA[lang][grp])

                if pattern:
                    data[lang][grp]["pattern"] = pattern(data[lang][grp]["data"])

    f = open(os.path.join(cur_dir, "data.py"), "w+")
    f.write(TEMPLATE)
    f.write(pprint.pformat(data, width=1))
    f.write(TEMPLATE_END)
    f.close()
<|endoftext|>


<commit_before># Copyright (c) 2018 Forschungszentrum Juelich GmbH
# Author: Yann Leprince <y.leprince@fz-juelich.de>
#
# This software is made available under the MIT licence, see LICENCE.txt.

"""Conversion of images to the Neuroglancer pre-computed format.

.. todo:: introduction to the high-level APIs
"""

__version__ = "0.1.0"
<commit_msg>Bump version for development<commit_after># Copyright (c) 2018 Forschungszentrum Juelich GmbH
# Author: Yann Leprince <y.leprince@fz-juelich.de>
#
# This software is made available under the MIT licence, see LICENCE.txt.

"""Conversion of images to the Neuroglancer pre-computed format.

.. todo:: introduction to the high-level APIs
"""

# Version used by setup.py and docs/conf.py
__version__ = "0.2.0.dev0"
<|endoftext|>


<commit_before>"""
@file coi-services/mi.idk/package_driver.py
@author Bill French
@brief Main script class for running the package_driver process
"""

import sys
import os.path
import zipfile

import yaml

from mi.core.log import get_logger ; log = get_logger()
from mi.idk import prompt
from mi.idk.metadata import Metadata
from mi.idk.nose_test import NoseTest
from mi.idk.driver_generator import DriverGenerator
from mi.idk.egg_generator import EggGenerator

class PackageManifest(object):
    """
    Object to create and store a package file manifest
    """

    ###
    #   Configuration
    ###
    def manifest_file(self):
        return "file.lst"

    def manifest_path(self):
        return "%s/%s" % (self.metadata.idk_dir(), self.manifest_file())

    ###
    #   Public Methods
    ###
    def __init__(self, metadata):
        """
        @brief ctor
        """
        self.metadata = metadata
        self.data = {}

    def add_file(self, source, description=None):
        """
        @brief Add a file to the file manifest
        @param source path the the file in the archive
        @description one line description of the file
        """

        if(not description): description = ''

        log.debug( "  ++ Adding " + source + " to manifest")
        self.data[source] = description

        self.save()

    def serialize(self):
        """
        @brief Serialize PackageManifest object data into a yaml string.
        @retval yaml string
        """
        return yaml.dump( self.data, default_flow_style=False )

    def save(self):
        """
        @brief Write YAML file with package manifest.
        """
        outputFile = self.manifest_path()

        if not os.path.exists(self.metadata.idk_dir()):
            os.makedirs(self.metadata.idk_dir())

        ofile = open( outputFile, 'w' )

        ofile.write( self.serialize() )
        ofile.close()

class PackageDriver(object):
    """
    Main class for running the package driver process.
    """

    ###
    #   Configuration
    ###
    def log_file(self):
        return "qualification.log"

    def log_path(self):
        return "%s/%s" % (self.metadata.idk_dir(), self.log_file())

    def archive_file(self):
        return "%s_%s_%s-%s-driver.zip" % (self.metadata.driver_make,
                                                      self.metadata.driver_model,
                                                      self.metadata.driver_name,
                                                      self.metadata.version)
    def archive_path(self):
        return os.path.join(os.path.expanduser("~"),self.archive_file())

    ###
    #   Public Methods
    ###
    def __init__(self):
        """
        @brief ctor
        """
        self.metadata = Metadata()
        self._zipfile = None
        self._manifest = None
        self._compression = None
        self.generator = DriverGenerator(self.metadata)

        # Set compression level
        self.zipfile_compression()

    def run_qualification_tests(self):
        """
        @brief Run all qualification tests for the driver and store the results for packaging
        """
        return True
        log.info("-- Running qualification tests")

        test = NoseTest(self.metadata, log_file=self.log_path())
        test.report_header()

        result = test.run_qualification()

        if(test.run_qualification()):
            log.info(" ++ Qualification tests passed")
            return True
        else:
            log.error("Qualification tests have fail!  No package created.")
            return False

    def package_driver(self):
        """
        @brief Store driver files in a zip package
        """
        log.info("-- Building driver package")
        self._store_package_files()

    def run(self):
        print "*** Starting Driver Packaging Process***"

        if( self.run_qualification_tests() ):
            self.package_driver()
            print "Package Created: " + self.archive_path()
        else:
            sys.exit()

    def zipfile(self):
        """
        @brief Return the ZipFile object.  Create the file if it isn't already open
        @retval ZipFile object
        """
        if(not self._zipfile):
            self._zipfile = zipfile.ZipFile(self.archive_path(), mode="w")

        return self._zipfile

    def zipfile_compression(self):
        """
        @brief What type of compression should we use for the package file.  If we have access to zlib, we will compress
        @retval Compression type
        """

        if(self._compression): return self._compression

        try:
            import zlib
            self._compression = zipfile.ZIP_DEFLATED
            log.info("Setting compression level to deflated")
        except:
            log.info("Setting compression level to store only")
            self._compression = zipfile.ZIP_STORED

    def manifest(self):
        """
        @brief Return the PackageManifest object.  Create it if it doesn't already exist
        @retval PackageManifest object
        """
        if(not self._manifest):
            self._manifest = PackageManifest(self.metadata)

        return self._manifest

    ###
    #   Private Methods
    ###
    def _store_package_files(self):
        """
        @brief Store all files in zip archive and add them to the manifest file
        """

        egg_generator = EggGenerator(self.metadata)
        egg_file = egg_generator.save()

        # Add egg
        self._add_file(egg_file, 'egg', 'python driver egg package')

        # Add the package metadata file
        self._add_file(self.metadata.metadata_path(), description = 'package metadata')

        # Add the qualification test log
        self._add_file(self.log_path(), description = 'qualification tests results')

        # Store additional resource files
        self._store_resource_files()

        # Finally save the manifest file.  This must be last of course
        self._add_file(self.manifest().manifest_path(), description = 'package manifest file')

    def _store_resource_files(self):
        """
        @brief Store additional files added by the driver developer.  These files life in the driver resource dir.
        """
        log.debug( " -- Searching for developer added resource files." )

        for file in os.listdir(self.generator.resource_dir()):
            log.debug("    ++ found: " + file)
            desc = prompt.text( 'Describe ' + file )
            self._add_file(self.generator.resource_dir() + "/" + file, 'resource', desc)

    def _add_file(self, source, destdir=None, description=None):
        """
        @brief Add a file to the zip package and store the file in the manifest.
        """
        filename = os.path.basename(source)
        dest = filename
        if(destdir):
            dest = "%s/%s" % (destdir, filename)

        log.debug( "archive %s to %s" % (filename, dest) )

        self.manifest().add_file(dest, description);
        self.zipfile().write(source, dest, self.zipfile_compression())

if __name__ == '__main__':
    app = PackageDriver()
    app.run()
<commit_msg>Added features to package_driver<commit_after>"""
@file coi-services/mi.idk/package_driver.py
@author Bill French
@brief Main script class for running the package_driver process
"""

import sys
import os.path
import zipfile

import yaml

from mi.core.log import get_logger ; log = get_logger()
from mi.idk import prompt
from mi.idk.metadata import Metadata
from mi.idk.nose_test import NoseTest
from mi.idk.driver_generator import DriverGenerator
from mi.idk.egg_generator import EggGenerator

class PackageManifest(object):
    """
    Object to create and store a package file manifest
    """

    ###
    #   Configuration
    ###
    def manifest_file(self):
        return "file.lst"

    def manifest_path(self):
        return "%s/%s" % (self.metadata.idk_dir(), self.manifest_file())

    ###
    #   Public Methods
    ###
    def __init__(self, metadata):
        """
        @brief ctor
        """
        self.metadata = metadata
        self.data = {}

    def add_file(self, source, description=None):
        """
        @brief Add a file to the file manifest
        @param source path the the file in the archive
        @description one line description of the file
        """

        if(not description): description = ''

        log.debug( "  ++ Adding " + source + " to manifest")
        self.data[source] = description

        self.save()

    def serialize(self):
        """
        @brief Serialize PackageManifest object data into a yaml string.
        @retval yaml string
        """
        return yaml.dump( self.data, default_flow_style=False )

    def save(self):
        """
        @brief Write YAML file with package manifest.
        """
        outputFile = self.manifest_path()

        if not os.path.exists(self.metadata.idk_dir()):
            os.makedirs(self.metadata.idk_dir())

        ofile = open( outputFile, 'w' )

        ofile.write( self.serialize() )
        ofile.close()

class PackageDriver(object):
    """
    Main class for running the package driver process.
    """

    ###
    #   Configuration
    ###
    def log_file(self):
        return "qualification.log"

    def log_path(self):
        return "%s/%s" % (self.metadata.idk_dir(), self.log_file())

    def archive_file(self):
        return "%s_%s_%s-%s-driver.zip" % (self.metadata.driver_make,
                                                      self.metadata.driver_model,
                                                      self.metadata.driver_name,
                                                      self.metadata.version)
    def archive_path(self):
        return os.path.join(os.path.expanduser("~"),self.archive_file())

    ###
    #   Public Methods
    ###
    def __init__(self):
        """
        @brief ctor
        """
        self.metadata = Metadata()
        self._zipfile = None
        self._manifest = None
        self._compression = None
        self.generator = DriverGenerator(self.metadata)

        # Set compression level
        self.zipfile_compression()

    def run_qualification_tests(self):
        """
        @brief Run all qualification tests for the driver and store the results for packaging
        """
        log.info("-- Running qualification tests")

        test = NoseTest(self.metadata, log_file=self.log_path())
        test.report_header()

        if(test.run_qualification()):
            log.info(" ++ Qualification tests passed")
            return True
        else:
            log.error("Qualification tests have fail!  No package created.")
            return False

    def package_driver(self):
        """
        @brief Store driver files in a zip package
        """
        log.info("-- Building driver package")
        self._store_package_files()

    def run(self):
        print "*** Starting Driver Packaging Process***"

        if len(sys.argv) == 2 and (sys.argv[1] == "--no-test"):
            # clear the log file so it exists
            f = open(self.log_path(), "w")
            f.write("Tests manually bypassed with --no-test option\n")
            f.close()
            self.package_driver()
        else:
            if(self.run_qualification_tests()):
                self.package_driver()

        print "Package Created: " + self.archive_path()

    def zipfile(self):
        """
        @brief Return the ZipFile object.  Create the file if it isn't already open
        @retval ZipFile object
        """
        if(not self._zipfile):
            self._zipfile = zipfile.ZipFile(self.archive_path(), mode="w")

        return self._zipfile

    def zipfile_compression(self):
        """
        @brief What type of compression should we use for the package file.  If we have access to zlib, we will compress
        @retval Compression type
        """

        if(self._compression): return self._compression

        try:
            import zlib
            self._compression = zipfile.ZIP_DEFLATED
            log.info("Setting compression level to deflated")
        except:
            log.info("Setting compression level to store only")
            self._compression = zipfile.ZIP_STORED

    def manifest(self):
        """
        @brief Return the PackageManifest object.  Create it if it doesn't already exist
        @retval PackageManifest object
        """
        if(not self._manifest):
            self._manifest = PackageManifest(self.metadata)

        return self._manifest

    ###
    #   Private Methods
    ###
    def _store_package_files(self):
        """
        @brief Store all files in zip archive and add them to the manifest file
        """

        egg_generator = EggGenerator(self.metadata)
        egg_file = egg_generator.save()

        # Add egg
        self._add_file(egg_file, 'egg', 'python driver egg package')

        # Add the package metadata file
        self._add_file(self.metadata.metadata_path(), description = 'package metadata')

        # Add the qualification test log
        self._add_file(self.log_path(), description = 'qualification tests results')

        # Store additional resource files
        self._store_resource_files()

        # Finally save the manifest file.  This must be last of course
        self._add_file(self.manifest().manifest_path(), description = 'package manifest file')

    def _store_resource_files(self):
        """
        @brief Store additional files added by the driver developer.  These
        files live in the driver resource dir.
        """
        log.debug(" -- Searching for developer added resource files.")
        resource_dir = self.generator.resource_dir()

        if os.path.exists(resource_dir):
            for file in os.listdir(resource_dir):
                log.debug("    ++ found: " + file)
                desc = prompt.text('Describe ' + file)
                self._add_file(resource_dir + "/" + file, 'resource', desc)
        else:
            log.debug(" --- No resource directory found, skipping...")

    def _add_file(self, source, destdir=None, description=None):
        """
        @brief Add a file to the zip package and store the file in the manifest.
        """
        filename = os.path.basename(source)
        dest = filename
        if(destdir):
            dest = "%s/%s" % (destdir, filename)

        log.debug("archive %s to %s" % (filename, dest))

        self.manifest().add_file(dest, description);
        self.zipfile().write(source, dest, self.zipfile_compression())

if __name__ == '__main__':
    app = PackageDriver()
    app.run()
<|endoftext|>


<commit_before>"""Defines top-level execution of test suite"""
# pylint: disable=missing-function-docstring

import os
from pathlib import Path

import nox

py_versions = ["3.8", "3.9"]

nox.options.sessions = [
    "flake8_diff",
    "flake8",
    "black",
    "pylint-3.8",
    "mypy-3.8",
    "unit_tests-3.8",
    "flake8_nb",
    "black_nb",
    "pylint_nb-3.8",
    "system_tests-3.8",
    "update_git_hooks",
]

# files we can run all the checks on, as they don't contain legacy code that
# has not yet been updated to pass all checks.
more_checks = [
    "metatlas/interfaces/compounds/populate.py",
    "metatlas/io/gdrive.py",
    "metatlas/io/file_converter.py",
    "metatlas/io/rclone.py",
    "metatlas/io/targeted_output.py",
    "metatlas/io/system_utils.py",
    "metatlas/io/write_utils.py",
    "metatlas/datastructures/atlas.py",
    "metatlas/datastructures/analysis_identifiers.py",
    "metatlas/datastructures/id_types.py",
    "metatlas/datastructures/metatlas_dataset.py",
    "metatlas/datastructures/spectrum.py",
    "metatlas/datastructures/utils.py",
    "metatlas/plots/compound_eic.py",
    "metatlas/plots/mirror.py",
    "metatlas/plots/plot_set.py",
    "metatlas/plots/tic.py",
    "metatlas/plots/utils.py",
    "metatlas/scripts/copy_ms_to_new_names.py",
    "metatlas/scripts/yaml_validation.py",
    "metatlas/targeted/rt_alignment.py",
    "metatlas/targeted/process.py",
    "metatlas/tools/add_msms_ref.py",
    "metatlas/tools/cheminfo.py",
    "metatlas/tools/config.py",
    "metatlas/tools/environment.py",
    "metatlas/tools/logging.py",
    "metatlas/tools/notebook.py",
    "metatlas/tools/parallel.py",
    "metatlas/tools/util.py",
    "metatlas/tools/validate_filenames.py",
    "noxfile.py",
    "tests",
]

# notebooks we can run all the checks on, as they don't contain legacy code that
# has not yet been updated to pass all checks.
notebooks = [
    "notebooks/reference/Targeted.ipynb",
    "notebooks/reference/RT_Alignment.ipynb",
    "notebooks/reference/Add_MSMS_Reference.ipynb",
]

pytest_deps = [
    "attrs==21.2.0",
    "coverage==6.0",
    "iniconfig==1.1.1",
    "numpy==1.22.4",
    "packaging==21.0",
    "pandas==1.4.2",
    "pluggy==0.13.1",
    "py==1.10.0",
    "pyparsing==2.4.7",
    "pytest==6.2.4",
    "pytest-cov==3.0.0",
    "pytest-mock==3.6.1",
    "toml==0.10.2",
]

mypy_deps = [
    "mypy==0.910",
    "types-PyYAML",
    "types-requests",
    "types-simplejson",
    "types-six",
    "types-tabulate",
]

pylint_deps = [
    "nox==2022.1.7",
    "pylint==2.11.1",
    "pytest==6.2.5",  # so "import pytest" doesn't get reported
]

nbqa_deps = [
    "nbqa==0.8.1",
    "tokenize-rt==4.1.0",
    "importlib-metadata==4.0.1",
    "astroid==2.8.0",
    "wrapt==1.12.1",
    "lazy_object_proxy==1.6.0",
    "isort==5.8.0",
]

flake8_deps = [
    "flake8==3.9.2",
    "flake8-bugbear==21.9.2",
    "flake8-builtins==1.5.3",
    "flake8-comprehensions==3.6.1",
]

pytest_flags = ["-vv", f"--basetemp={Path.home() / '.pytest_tmp'}"]

nox.options.error_on_external_run = True
REUSE_LARGE_VENV = True
NB_LINE_LEN = 140

# parallel testings doesn't work well on NERSC login nodes
if "NERSC_HOST" not in os.environ:
    pytest_deps.append("pytest-xdist[psutil]==2.4.0")
    pytest_flags.extend(["--numprocesses", "auto"])

@nox.session(python=py_versions[0])
def flake8_diff(session):
    session.install(*flake8_deps)
    session.run("sh", "-c", "git diff -U0 -w --staged HEAD | flake8 --diff", external=True)

@nox.session(python=py_versions[0])
def flake8_all(session):
    session.install(*flake8_deps)
    session.run("flake8", "metatlas", "tests")

@nox.session(python=py_versions[0])
def flake8(session):
    session.install(*flake8_deps)
    session.run("flake8", *more_checks)

@nox.session(python=py_versions[0])
def black_all(session):
    session.install("black")
    session.run("black", "--check", "--diff", "--color", "metatlas", "tests")

@nox.session(python=py_versions[0])
def black(session):
    session.install("black")
    session.run("black", "--check", "--diff", "--color", *more_checks)

@nox.session(python=py_versions[0])
def blacken(session):
    """this modifies the files to meet black's requirements"""
    session.install("black")
    session.run("black", *more_checks)

@nox.session(python=py_versions, reuse_venv=REUSE_LARGE_VENV)
def mypy(session):
    session.install("-r", "docker/requirements.txt", *mypy_deps)
    session.run("mypy", "metatlas")

@nox.session(python=py_versions, reuse_venv=REUSE_LARGE_VENV)
def pylint(session):
    session.install("-r", "docker/requirements.txt", *pylint_deps)
    session.run("pylint", *more_checks)

@nox.session(python=py_versions, reuse_venv=REUSE_LARGE_VENV)
def pylint_nb(session):
    session.install("-r", "docker/requirements.txt", *nbqa_deps, *pylint_deps)
    # dupliate code cannot be disabled on per-cell level https://github.com/PyCQA/pylint/issues/214
    # Some duplicate code is required to setup the notebook and do error handling.
    # So turn off duplicate code for whole session -- not ideal.
    session.run("nbqa", "pylint", "--disable=duplicate-code", f"--max-line-length={NB_LINE_LEN}", *notebooks)

@nox.session(python=py_versions[0])
def flake8_nb(session):
    session.install(*nbqa_deps, *flake8_deps)
    session.run("nbqa", "flake8", *notebooks)

@nox.session(python=py_versions[0])
def black_nb(session):
    session.install("black", *nbqa_deps)
    session.run("nbqa", "black", f"--line-length={NB_LINE_LEN}", "--check", *notebooks)

@nox.session(python=py_versions[0])
def blacken_nb(session):
    """this modifies notebook files to meet black's requirements"""
    session.install("black", *nbqa_deps)
    session.run("nbqa", "black", f"--line-length={NB_LINE_LEN}", "--nbqa-mutate", *notebooks)

@nox.session(python=py_versions, reuse_venv=REUSE_LARGE_VENV)
def unit_tests(session):
    session.install("-r", "docker/requirements.txt", *pytest_deps)
    session.run(
        "pytest",
        *pytest_flags,
        *session.posargs,
        "--cov",
        "metatlas",
        "--cov-report",
        "term:skip-covered",
        "--cov-branch",
        "tests/unit/",
        env={"METATLAS_LOCAL": "TRUE"},
    )

@nox.session(python=py_versions[0], reuse_venv=REUSE_LARGE_VENV)
def cov_report(session):
    session.install("-r", "docker/requirements.txt", *pytest_deps)
    session.run(
        "pytest",
        *session.posargs,
        "--cov",
        "metatlas",
        "--cov-report",
        "term-missing:skip-covered",
        "--cov-branch",
        "tests/unit/",
        env={"METATLAS_LOCAL": "TRUE"},
    )

@nox.session(python=py_versions[0])
def system_tests(session):
    session.install(*pytest_deps)
    session.run("pytest", *pytest_flags, *session.posargs, "tests/system/")

@nox.session(python=py_versions[0])
def install_git_hooks(session):
    session.install("pre-commit")
    session.run("pre-commit", "install")

@nox.session(python=py_versions[0])
def update_git_hooks(session):
    session.install("pre-commit")
    session.run("pre-commit", "autoupdate")
<commit_msg>upgrade nbqa to v1.4.0<commit_after>"""Defines top-level execution of test suite"""
# pylint: disable=missing-function-docstring

import os
from pathlib import Path

import nox

py_versions = ["3.8", "3.9"]

nox.options.sessions = [
    "flake8_diff",
    "flake8",
    "black",
    "pylint-3.8",
    "mypy-3.8",
    "unit_tests-3.8",
    "flake8_nb",
    "black_nb",
    "pylint_nb-3.8",
    "system_tests-3.8",
    "update_git_hooks",
]

# files we can run all the checks on, as they don't contain legacy code that
# has not yet been updated to pass all checks.
more_checks = [
    "metatlas/interfaces/compounds/populate.py",
    "metatlas/io/gdrive.py",
    "metatlas/io/file_converter.py",
    "metatlas/io/rclone.py",
    "metatlas/io/targeted_output.py",
    "metatlas/io/system_utils.py",
    "metatlas/io/write_utils.py",
    "metatlas/datastructures/atlas.py",
    "metatlas/datastructures/analysis_identifiers.py",
    "metatlas/datastructures/id_types.py",
    "metatlas/datastructures/metatlas_dataset.py",
    "metatlas/datastructures/spectrum.py",
    "metatlas/datastructures/utils.py",
    "metatlas/plots/compound_eic.py",
    "metatlas/plots/mirror.py",
    "metatlas/plots/plot_set.py",
    "metatlas/plots/tic.py",
    "metatlas/plots/utils.py",
    "metatlas/scripts/copy_ms_to_new_names.py",
    "metatlas/scripts/yaml_validation.py",
    "metatlas/targeted/rt_alignment.py",
    "metatlas/targeted/process.py",
    "metatlas/tools/add_msms_ref.py",
    "metatlas/tools/cheminfo.py",
    "metatlas/tools/config.py",
    "metatlas/tools/environment.py",
    "metatlas/tools/logging.py",
    "metatlas/tools/notebook.py",
    "metatlas/tools/parallel.py",
    "metatlas/tools/util.py",
    "metatlas/tools/validate_filenames.py",
    "noxfile.py",
    "tests",
]

# notebooks we can run all the checks on, as they don't contain legacy code that
# has not yet been updated to pass all checks.
notebooks = [
    "notebooks/reference/Targeted.ipynb",
    "notebooks/reference/RT_Alignment.ipynb",
    "notebooks/reference/Add_MSMS_Reference.ipynb",
]

pytest_deps = [
    "attrs==21.2.0",
    "coverage==6.0",
    "iniconfig==1.1.1",
    "numpy==1.22.4",
    "packaging==21.0",
    "pandas==1.4.2",
    "pluggy==0.13.1",
    "py==1.10.0",
    "pyparsing==2.4.7",
    "pytest==6.2.4",
    "pytest-cov==3.0.0",
    "pytest-mock==3.6.1",
    "toml==0.10.2",
]

mypy_deps = [
    "mypy==0.910",
    "types-PyYAML",
    "types-requests",
    "types-simplejson",
    "types-six",
    "types-tabulate",
]

pylint_deps = [
    "nox==2022.1.7",
    "pylint==2.11.1",
    "pytest==6.2.5",  # so "import pytest" doesn't get reported
]

nbqa_deps = [
    "nbqa==1.4.0",
    "tokenize-rt==4.1.0",
    "importlib-metadata==4.0.1",
    "astroid==2.8.0",
    "wrapt==1.12.1",
    "lazy_object_proxy==1.6.0",
    "isort==5.8.0",
]

flake8_deps = [
    "flake8==3.9.2",
    "flake8-bugbear==21.9.2",
    "flake8-builtins==1.5.3",
    "flake8-comprehensions==3.6.1",
]

pytest_flags = ["-vv", f"--basetemp={Path.home() / '.pytest_tmp'}"]
nbqa_flags = ["--nbqa-dont-skip-bad-cells"]

nox.options.error_on_external_run = True
REUSE_LARGE_VENV = True
NB_LINE_LEN = 140

# parallel testings doesn't work well on NERSC login nodes
if "NERSC_HOST" not in os.environ:
    pytest_deps.append("pytest-xdist[psutil]==2.4.0")
    pytest_flags.extend(["--numprocesses", "auto"])

@nox.session(python=py_versions[0])
def flake8_diff(session):
    session.install(*flake8_deps)
    session.run("sh", "-c", "git diff -U0 -w --staged HEAD | flake8 --diff", external=True)

@nox.session(python=py_versions[0])
def flake8_all(session):
    session.install(*flake8_deps)
    session.run("flake8", "metatlas", "tests")

@nox.session(python=py_versions[0])
def flake8(session):
    session.install(*flake8_deps)
    session.run("flake8", *more_checks)

@nox.session(python=py_versions[0])
def black_all(session):
    session.install("black")
    session.run("black", "--check", "--diff", "--color", "metatlas", "tests")

@nox.session(python=py_versions[0])
def black(session):
    session.install("black")
    session.run("black", "--check", "--diff", "--color", *more_checks)

@nox.session(python=py_versions[0])
def blacken(session):
    """this modifies the files to meet black's requirements"""
    session.install("black")
    session.run("black", *more_checks)

@nox.session(python=py_versions, reuse_venv=REUSE_LARGE_VENV)
def mypy(session):
    session.install("-r", "docker/requirements.txt", *mypy_deps)
    session.run("mypy", "metatlas")

@nox.session(python=py_versions, reuse_venv=REUSE_LARGE_VENV)
def pylint(session):
    session.install("-r", "docker/requirements.txt", *pylint_deps)
    session.run("pylint", *more_checks)

@nox.session(python=py_versions, reuse_venv=REUSE_LARGE_VENV)
def pylint_nb(session):
    session.install("-r", "docker/requirements.txt", *nbqa_deps, *pylint_deps)
    # dupliate code cannot be disabled on per-cell level https://github.com/PyCQA/pylint/issues/214
    # Some duplicate code is required to setup the notebook and do error handling.
    # So turn off duplicate code for whole session -- not ideal.
    session.run("nbqa", "pylint", "--disable=duplicate-code", f"--max-line-length={NB_LINE_LEN}", *nbqa_flags, *notebooks)

@nox.session(python=py_versions[0])
def flake8_nb(session):
    session.install(*nbqa_deps, *flake8_deps)
    session.run("nbqa", "flake8", *nbqa_flags, *notebooks)

@nox.session(python=py_versions[0])
def black_nb(session):
    session.install("black", *nbqa_deps)
    session.run("nbqa", "black", f"--line-length={NB_LINE_LEN}", "--check", *nbqa_flags, *notebooks)

@nox.session(python=py_versions[0])
def blacken_nb(session):
    """this modifies notebook files to meet black's requirements"""
    session.install("black", *nbqa_deps)
    session.run("nbqa", "black", f"--line-length={NB_LINE_LEN}", *nbqa_flags, *notebooks)

@nox.session(python=py_versions, reuse_venv=REUSE_LARGE_VENV)
def unit_tests(session):
    session.install("-r", "docker/requirements.txt", *pytest_deps)
    session.run(
        "pytest",
        *pytest_flags,
        *session.posargs,
        "--cov",
        "metatlas",
        "--cov-report",
        "term:skip-covered",
        "--cov-branch",
        "tests/unit/",
        env={"METATLAS_LOCAL": "TRUE"},
    )

@nox.session(python=py_versions[0], reuse_venv=REUSE_LARGE_VENV)
def cov_report(session):
    session.install("-r", "docker/requirements.txt", *pytest_deps)
    session.run(
        "pytest",
        *session.posargs,
        "--cov",
        "metatlas",
        "--cov-report",
        "term-missing:skip-covered",
        "--cov-branch",
        "tests/unit/",
        env={"METATLAS_LOCAL": "TRUE"},
    )

@nox.session(python=py_versions[0])
def system_tests(session):
    session.install(*pytest_deps)
    session.run("pytest", *pytest_flags, *session.posargs, "tests/system/")

@nox.session(python=py_versions[0])
def install_git_hooks(session):
    session.install("pre-commit")
    session.run("pre-commit", "install")

@nox.session(python=py_versions[0])
def update_git_hooks(session):
    session.install("pre-commit")
    session.run("pre-commit", "autoupdate")
<|endoftext|>


<commit_before>#!/usr/bin/python

import os,sys
#from optparse import OptionParser
import argparse

from biomaj.bank import Bank
from biomaj.config import BiomajConfig
from biomaj.notify import Notify
from biomaj.options import Options

def main():

  #parser = OptionParser()
  #parser.add_option('-c', '--config', dest="config",help="Configuration file")
  #parser.add_option('-u', '--update', dest="update", help="Update action", action="store_true", default=False)
  #parser.add_option('--fromscratch', dest="fromscratch", help="Force a new cycle update", action="store_true", default=False)
  #parser.add_option('--publish', dest="publish", help="Publish", action="store_true", default=False)
  #parser.add_option('--release', dest="release",help="release of the bank")
  #parser.add_option('-r', '--remove', dest="remove", help="Update action", action="store_true", default=False)
  #parser.add_option('-s', '--status', dest="status", help="Get status", action="store_true", default=False)
  #parser.add_option('-b', '--bank', dest="bank",help="bank name")
  #(options, args) = parser.parse_args()

  parser = argparse.ArgumentParser()
  parser.add_argument('-c', '--config', dest="config",help="Configuration file")
  parser.add_argument('-u', '--update', dest="update", help="Update action", action="store_true", default=False)
  parser.add_argument('--fromscratch', dest="fromscratch", help="Force a new cycle update", action="store_true", default=False)
  parser.add_argument('--publish', dest="publish", help="Publish", action="store_true", default=False)
  parser.add_argument('--release', dest="release",help="release of the bank")
  parser.add_argument('--from-task', dest="from_task",help="Start cycle at a specific task (init always executed)")
  parser.add_argument('--process', dest="process",help="Linked to from-task, optionally specify a block, meta or process name to start from")
  parser.add_argument('--log', dest="log",help="log level")
  parser.add_argument('-r', '--remove', dest="remove", help="Remove a bank release", action="store_true", default=False)
  parser.add_argument('--remove-all', dest="removeall", help="Remove all bank releases and database records", action="store_true", default=False)
  parser.add_argument('-s', '--status', dest="status", help="Get status", action="store_true", default=False)
  parser.add_argument('-b', '--bank', dest="bank",help="bank name")
  parser.add_argument('--stop-before', dest="stop_before",help="Store workflow before task")
  parser.add_argument('--stop-after', dest="stop_after",help="Store workflow after task")
  parser.add_argument('--freeze', dest="freeze", help="Freeze a bank release", action="store_true", default=False)
  parser.add_argument('--unfreeze', dest="unfreeze", help="Unfreeze a bank release", action="store_true", default=False)
  parser.add_argument('--force', dest="force", help="Force action", action="store_true", default=False)

  options = Options()
  parser.parse_args(namespace=options)

  bmaj = None
  if options.config is not None:
    BiomajConfig.load_config(options.config)

  if options.status:
    if options.bank:
      bank = Bank(options.bank)
      _bank = bank.bank
      print '#' * 80
      print "# Name:\t"+_bank['name']
      print "# Type:\t"+_bank['properties']['type']
      release = None
      if 'current' in _bank and _bank['current']:
        for prod in _bank['production']:
          if _bank['current'] == prod['session']:
            release = prod['release']
      print "# Published release:\t"+str(release)
      print "# Production directories"
      for prod in _bank['production']:
        if 'freeze' in prod:
          print "#\tFreeze:\t"+prod['freeze']
        print "#\tRelease:\t"+prod['release']
        print "#\t\tSession:\t"+str(prod['session'])
        release_dir = os.path.join(bank.config.get('data.dir'),
                      bank.config.get('dir.version'),
                      prod['prod_dir'])
        print "#\t\tDirectory:\t"+release_dir
      print '#' * 80
    else:
      print '#' * 80
      print "# Name\tType\tRelease"
      banks = Bank.list()
      for bank in banks:
        '''
        production = { 'release': self.session.get('release'),
                        'session': self.session._session['id'],
                        'data_dir': self.config.get('data.dir'),
                        'prod_dir': self.session.get_release_directory()}
        '''
        if 'current' in bank and bank['current']:
          for prod in bank['production']:
            if bank['current'] == prod['session']:
              release = prod['release']
        else:
          release = None
        print " "+bank['name']+"\t"+bank['properties']['type']+"\t"+str(release)
      print '#' * 80
      return

  if options.update and options.bank:
    bmaj = Bank(options.bank, options)
    print 'Log file: '+bmaj.config.log_file
    res = bmaj.update()
    Notify.notifyBankAction(bmaj)
    if not res:
      sys.exit(1)

  if options.freeze and options.release and options.bank:
    bmaj = Bank(options.bank, options)
    res = bmaj.freeze(options.release)
    if not res:
      sys.exit(1)

  if options.unfreeze and options.release and options.bank:
    bmaj = Bank(options.bank, options)
    res = bmaj.unfreeze(options.release)
    if not res:
      sys.exit(1)

  if ((options.remove and options.release) or options.removeall) and options.bank:
    bmaj = Bank(options.bank, options)
    print 'Log file: '+bmaj.config.log_file
    if options.removeall:
      res = bmaj.removeAll(options.force)
    else:
      res = bmaj.remove(options.release)
      Notify.notifyBankAction(bmaj)
    if not res:
      sys.exit(1)

  if options.publish:
    if not options.bank:
      print "Bank name or release is missing"
      sys.exit(1)
    bmaj = Bank(options.bank, options)
    print 'Log file: '+bmaj.config.log_file
    bmaj.load_session()
    bank = bmaj.bank
    session = None
    if options.get_option('release') is None:
      # Get latest prod release
      if len(bank['production'])>0:
        prod = bank['production'][len(bank['production'])-1]
        for s in bank['sessions']:
          if s['id'] == prod['session']:
            session = s
            break
    else:
      # Search production release matching release
      for prod in bank['production']:
        if prod['release'] == options.release or prod['prod_dir'] == options.release:
          # Search session related to this production release
          for s in bank['sessions']:
            if s['id'] == prod['session']:
              session = s
              break
          break
    if session is None:
      print "No production session could be found for this release"
      sys.exit(1)
    bmaj.session._session = session
    bmaj.publish()

if __name__ == '__main__':
    main()
<commit_msg>fix print error<commit_after>#!/usr/bin/python

import os,sys
#from optparse import OptionParser
import argparse

from biomaj.bank import Bank
from biomaj.config import BiomajConfig
from biomaj.notify import Notify
from biomaj.options import Options

def main():

  #parser = OptionParser()
  #parser.add_option('-c', '--config', dest="config",help="Configuration file")
  #parser.add_option('-u', '--update', dest="update", help="Update action", action="store_true", default=False)
  #parser.add_option('--fromscratch', dest="fromscratch", help="Force a new cycle update", action="store_true", default=False)
  #parser.add_option('--publish', dest="publish", help="Publish", action="store_true", default=False)
  #parser.add_option('--release', dest="release",help="release of the bank")
  #parser.add_option('-r', '--remove', dest="remove", help="Update action", action="store_true", default=False)
  #parser.add_option('-s', '--status', dest="status", help="Get status", action="store_true", default=False)
  #parser.add_option('-b', '--bank', dest="bank",help="bank name")
  #(options, args) = parser.parse_args()

  parser = argparse.ArgumentParser()
  parser.add_argument('-c', '--config', dest="config",help="Configuration file")
  parser.add_argument('-u', '--update', dest="update", help="Update action", action="store_true", default=False)
  parser.add_argument('--fromscratch', dest="fromscratch", help="Force a new cycle update", action="store_true", default=False)
  parser.add_argument('--publish', dest="publish", help="Publish", action="store_true", default=False)
  parser.add_argument('--release', dest="release",help="release of the bank")
  parser.add_argument('--from-task', dest="from_task",help="Start cycle at a specific task (init always executed)")
  parser.add_argument('--process', dest="process",help="Linked to from-task, optionally specify a block, meta or process name to start from")
  parser.add_argument('--log', dest="log",help="log level")
  parser.add_argument('-r', '--remove', dest="remove", help="Remove a bank release", action="store_true", default=False)
  parser.add_argument('--remove-all', dest="removeall", help="Remove all bank releases and database records", action="store_true", default=False)
  parser.add_argument('-s', '--status', dest="status", help="Get status", action="store_true", default=False)
  parser.add_argument('-b', '--bank', dest="bank",help="bank name")
  parser.add_argument('--stop-before', dest="stop_before",help="Store workflow before task")
  parser.add_argument('--stop-after', dest="stop_after",help="Store workflow after task")
  parser.add_argument('--freeze', dest="freeze", help="Freeze a bank release", action="store_true", default=False)
  parser.add_argument('--unfreeze', dest="unfreeze", help="Unfreeze a bank release", action="store_true", default=False)
  parser.add_argument('--force', dest="force", help="Force action", action="store_true", default=False)

  options = Options()
  parser.parse_args(namespace=options)

  bmaj = None
  if options.config is not None:
    BiomajConfig.load_config(options.config)

  if options.status:
    if options.bank:
      bank = Bank(options.bank)
      _bank = bank.bank
      print '#' * 80
      print "# Name:\t"+_bank['name']
      print "# Type:\t"+_bank['properties']['type']
      release = None
      if 'current' in _bank and _bank['current']:
        for prod in _bank['production']:
          if _bank['current'] == prod['session']:
            release = prod['release']
      print "# Published release:\t"+str(release)
      print "# Production directories"
      for prod in _bank['production']:
        if 'freeze' in prod:
          print "#\tFreeze:\t"+str(prod['freeze'])
        print "#\tRelease:\t"+prod['release']
        print "#\t\tSession:\t"+str(prod['session'])
        release_dir = os.path.join(bank.config.get('data.dir'),
                      bank.config.get('dir.version'),
                      prod['prod_dir'])
        print "#\t\tDirectory:\t"+release_dir
      print '#' * 80
    else:
      print '#' * 80
      print "# Name\tType\tRelease"
      banks = Bank.list()
      for bank in banks:
        '''
        production = { 'release': self.session.get('release'),
                        'session': self.session._session['id'],
                        'data_dir': self.config.get('data.dir'),
                        'prod_dir': self.session.get_release_directory()}
        '''
        if 'current' in bank and bank['current']:
          for prod in bank['production']:
            if bank['current'] == prod['session']:
              release = prod['release']
        else:
          release = None
        print " "+bank['name']+"\t"+bank['properties']['type']+"\t"+str(release)
      print '#' * 80
      return

  if options.update and options.bank:
    bmaj = Bank(options.bank, options)
    print 'Log file: '+bmaj.config.log_file
    res = bmaj.update()
    Notify.notifyBankAction(bmaj)
    if not res:
      sys.exit(1)

  if options.freeze and options.release and options.bank:
    bmaj = Bank(options.bank, options)
    res = bmaj.freeze(options.release)
    if not res:
      sys.exit(1)

  if options.unfreeze and options.release and options.bank:
    bmaj = Bank(options.bank, options)
    res = bmaj.unfreeze(options.release)
    if not res:
      sys.exit(1)

  if ((options.remove and options.release) or options.removeall) and options.bank:
    bmaj = Bank(options.bank, options)
    print 'Log file: '+bmaj.config.log_file
    if options.removeall:
      res = bmaj.removeAll(options.force)
    else:
      res = bmaj.remove(options.release)
      Notify.notifyBankAction(bmaj)
    if not res:
      sys.exit(1)

  if options.publish:
    if not options.bank:
      print "Bank name or release is missing"
      sys.exit(1)
    bmaj = Bank(options.bank, options)
    print 'Log file: '+bmaj.config.log_file
    bmaj.load_session()
    bank = bmaj.bank
    session = None
    if options.get_option('release') is None:
      # Get latest prod release
      if len(bank['production'])>0:
        prod = bank['production'][len(bank['production'])-1]
        for s in bank['sessions']:
          if s['id'] == prod['session']:
            session = s
            break
    else:
      # Search production release matching release
      for prod in bank['production']:
        if prod['release'] == options.release or prod['prod_dir'] == options.release:
          # Search session related to this production release
          for s in bank['sessions']:
            if s['id'] == prod['session']:
              session = s
              break
          break
    if session is None:
      print "No production session could be found for this release"
      sys.exit(1)
    bmaj.session._session = session
    bmaj.publish()

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>import hashlib
import os

import six
import libcloud.security
from libcloud.storage.providers import get_driver
from libcloud.storage.types import ContainerDoesNotExistError, ObjectDoesNotExistError
from six.moves.urllib.parse import urlparse

# Include the current cURL CA bundle as a fallback
_base_path = os.path.abspath(os.path.dirname(__file__))
libcloud.security.CA_CERTS_PATH.append(os.path.join(_base_path, 'cacert.pem'))

class Sizer(object):
    def __init__(self):
        self.size = 0

    def update(self, data):
        self.size += len(data)

class StorageWrapper(object):
    def __init__(self, uri, no_public=False):
        self.uri = urlparse(uri)
        self.no_public = no_public
        self.storage = self._get_storage(self.uri)
        self._hashes = {}

    def upload(self, path, data):
        if isinstance(data, six.binary_type):
            # Simple case of in-memory content
            return self.upload(path, (data,))
        elif hasattr(data, 'read'):
            # File-like object
            def it(size=4096):
                buf = data.read(size)
                while buf:
                    yield buf
                    buf = data.read(size)
            return self.upload(path, it())
        else:
            # Just try iterating
            self._update_hashes(path, '')
            path_root, path_ext = os.path.splitext(path)
            extra = {}
            if path_ext == '.gz':
                extra['content_type'] = 'application/x-gzip'
            else:
                extra['content_type'] = 'text/plain'
            if not self.no_public:
                if self.uri.scheme.startswith('s3'):
                    # Yeah so this doesn't work, scroll down to the bottom and cry
                    extra['acl'] = 'public-read'
                    extra['meta_data'] = {'acl': 'public-read'}
                # Other clouds here
            return self.storage.upload_object_via_stream((self._update_hashes(path, buf, False) for buf in data), path, extra=extra)

    def download(self, path, skip_hash=False):
        # Assumption, this isn't a big file
        try:
            obj = self.storage.get_object(path)
        except ObjectDoesNotExistError:
            return None
        if not skip_hash:
            self._update_hashes(path, '')
        return ''.join(
            buf if skip_hash else self._update_hashes(path, buf, False)
            for buf in self.storage.download_object_as_stream(obj)
        )

    def __contains__(self, path):
        try:
            self.storage.get_object(path)
            return True
        except ObjectDoesNotExistError:
            return False

    def _update_hashes(self, path, data, reset=True):
        if reset or path not in self._hashes:
            self._hashes[path] = {
                'md5': hashlib.md5(),
                'sha1': hashlib.sha1(),
                'sha256': hashlib.sha256(),
                'size': Sizer(),
            }
        if data:
            for hasher in six.itervalues(self._hashes[path]):
                hasher.update(data)
        return data

    def hashes(self, path):
        if path not in self._hashes:
            self._update_hashes(path, '')
        return self._hashes[path]

    @classmethod
    def _get_storage(cls, uri):
        """
        Given a URI like local:///srv/repo or s3://key:secret@apt.example.com,
        return a libcloud storage or container object.
        """
        driver = cls._get_driver(uri.scheme)
        key = uri.username
        secret = uri.password
        container = uri.netloc
        driver_kwargs = {}
        if uri.scheme.startswith('s3'):
            if not key:
                key = os.environ.get('AWS_ACCESS_KEY_ID')
            if not secret:
                secret = os.environ.get('AWS_SECRET_ACCESS_KEY')
            if not (key and secret and container):
                raise ValueError('For S3 you must provide an access key ID, secret access key, and bucket name')
            # No way to store this in the URI, what about a CLI option too?
            if 'AWS_TOKEN' in os.environ:
                driver_kwargs['token'] = os.environ['AWS_TOKEN']
        elif uri.scheme == 'local':
            parts = []
            if uri.netloc:
                parts.append(uri.netloc)
            if uri.path:
                parts.append(uri.path)
            if not parts:
                parts.append('.')
            base_path = os.path.abspath(''.join(parts))
            key = os.path.dirname(base_path)
            container = os.path.basename(base_path)
        storage = driver(key, secret, **driver_kwargs)
        try:
            return storage.get_container(container)
        except ContainerDoesNotExistError:
            return storage.create_container(container)

    @classmethod
    def _get_driver(cls, name):
        """Wrapper for libcloud's get_driver for testing."""
        return get_driver(name)

# OH GOD I DON'T EVEN
from libcloud.storage.drivers.s3 import S3StorageDriver
old_upload_object = S3StorageDriver._upload_object

def _upload_object(*args, **kwargs):
    # Rewrite to the correct header
    if kwargs.get('headers', {}).get('x-amz-meta-acl'):
        kwargs['headers']['x-amz-acl'] = kwargs['headers'].pop('x-amz-meta-acl')
    return old_upload_object(*args, **kwargs)
S3StorageDriver._upload_object = _upload_object
<commit_msg>Expose a download_iter to get an iterable instead of a joined file. <commit_after>import hashlib
import os

import six
import libcloud.security
from libcloud.storage.providers import get_driver
from libcloud.storage.types import ContainerDoesNotExistError, ObjectDoesNotExistError
from six.moves.urllib.parse import urlparse

# Include the current cURL CA bundle as a fallback
_base_path = os.path.abspath(os.path.dirname(__file__))
libcloud.security.CA_CERTS_PATH.append(os.path.join(_base_path, 'cacert.pem'))

class Sizer(object):
    def __init__(self):
        self.size = 0

    def update(self, data):
        self.size += len(data)

class StorageWrapper(object):
    def __init__(self, uri, no_public=False):
        self.uri = urlparse(uri)
        self.no_public = no_public
        self.storage = self._get_storage(self.uri)
        self._hashes = {}

    def upload(self, path, data):
        if isinstance(data, six.binary_type):
            # Simple case of in-memory content
            return self.upload(path, (data,))
        elif hasattr(data, 'read'):
            # File-like object
            def it(size=4096):
                buf = data.read(size)
                while buf:
                    yield buf
                    buf = data.read(size)
            return self.upload(path, it())
        else:
            # Just try iterating
            self._update_hashes(path, '')
            path_root, path_ext = os.path.splitext(path)
            extra = {}
            if path_ext == '.gz':
                extra['content_type'] = 'application/x-gzip'
            else:
                extra['content_type'] = 'text/plain'
            if not self.no_public:
                if self.uri.scheme.startswith('s3'):
                    # Yeah so this doesn't work, scroll down to the bottom and cry
                    extra['acl'] = 'public-read'
                    extra['meta_data'] = {'acl': 'public-read'}
                # Other clouds here
            return self.storage.upload_object_via_stream((self._update_hashes(path, buf, False) for buf in data), path, extra=extra)

    def download(self, path, skip_hash=False):
        # Assumption, this isn't a big file
        it = self.download_iter(path, skip_hash)
        if it:
            return ''.join(it)

    def download_iter(self, path, skip_hash=False):
        try:
            obj = self.storage.get_object(path)
        except ObjectDoesNotExistError:
            return None
        if not skip_hash:
            self._update_hashes(path, '')
        return (
            buf if skip_hash else self._update_hashes(path, buf, False)
            for buf in self.storage.download_object_as_stream(obj)
        )

    def __contains__(self, path):
        try:
            self.storage.get_object(path)
            return True
        except ObjectDoesNotExistError:
            return False

    def _update_hashes(self, path, data, reset=True):
        if reset or path not in self._hashes:
            self._hashes[path] = {
                'md5': hashlib.md5(),
                'sha1': hashlib.sha1(),
                'sha256': hashlib.sha256(),
                'size': Sizer(),
            }
        if data:
            for hasher in six.itervalues(self._hashes[path]):
                hasher.update(data)
        return data

    def hashes(self, path):
        if path not in self._hashes:
            self._update_hashes(path, '')
        return self._hashes[path]

    @classmethod
    def _get_storage(cls, uri):
        """
        Given a URI like local:///srv/repo or s3://key:secret@apt.example.com,
        return a libcloud storage or container object.
        """
        driver = cls._get_driver(uri.scheme)
        key = uri.username
        secret = uri.password
        container = uri.netloc
        driver_kwargs = {}
        if uri.scheme.startswith('s3'):
            if not key:
                key = os.environ.get('AWS_ACCESS_KEY_ID')
            if not secret:
                secret = os.environ.get('AWS_SECRET_ACCESS_KEY')
            if not (key and secret and container):
                raise ValueError('For S3 you must provide an access key ID, secret access key, and bucket name')
            # No way to store this in the URI, what about a CLI option too?
            if 'AWS_TOKEN' in os.environ:
                driver_kwargs['token'] = os.environ['AWS_TOKEN']
        elif uri.scheme == 'local':
            parts = []
            if uri.netloc:
                parts.append(uri.netloc)
            if uri.path:
                parts.append(uri.path)
            if not parts:
                parts.append('.')
            base_path = os.path.abspath(''.join(parts))
            key = os.path.dirname(base_path)
            container = os.path.basename(base_path)
        storage = driver(key, secret, **driver_kwargs)
        try:
            return storage.get_container(container)
        except ContainerDoesNotExistError:
            return storage.create_container(container)

    @classmethod
    def _get_driver(cls, name):
        """Wrapper for libcloud's get_driver for testing."""
        return get_driver(name)

# OH GOD I DON'T EVEN
from libcloud.storage.drivers.s3 import S3StorageDriver
old_upload_object = S3StorageDriver._upload_object

def _upload_object(*args, **kwargs):
    # Rewrite to the correct header
    if kwargs.get('headers', {}).get('x-amz-meta-acl'):
        kwargs['headers']['x-amz-acl'] = kwargs['headers'].pop('x-amz-meta-acl')
    return old_upload_object(*args, **kwargs)
S3StorageDriver._upload_object = _upload_object
<|endoftext|>