text: #
  # Licensed to the Apache Software Foundation (ASF) under one or more
  # contributor license agreements.  See the NOTICE file distributed with
  # this work for additional information regarding copyright ownership.
  # The ASF licenses this file to You under the Apache License, Version 2.0
  # (the "License"); you may not use this file except in compliance with
  # the License.  You may obtain a copy of the License at
  #
  #    http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  
  """Support for Dataflow triggers.
  
  Triggers control when in processing time windows get emitted.
  """
  
  # pytype: skip-file
  
  from __future__ import absolute_import
  
  import collections
  import copy
  import logging
  import numbers
  from abc import ABCMeta
  from abc import abstractmethod
  from builtins import object
  
  from future.moves.itertools import zip_longest
  from future.utils import iteritems
  from future.utils import with_metaclass
  
  from apache_beam.coders import coder_impl
  from apache_beam.coders import observable
  from apache_beam.portability.api import beam_runner_api_pb2
  from apache_beam.transforms import combiners
  from apache_beam.transforms import core
  from apache_beam.transforms.timeutil import TimeDomain
  from apache_beam.transforms.window import GlobalWindow
  from apache_beam.transforms.window import GlobalWindows
  from apache_beam.transforms.window import TimestampCombiner
  from apache_beam.transforms.window import WindowedValue
  from apache_beam.transforms.window import WindowFn
  from apache_beam.utils import windowed_value
  from apache_beam.utils.timestamp import MAX_TIMESTAMP
  from apache_beam.utils.timestamp import MIN_TIMESTAMP
  from apache_beam.utils.timestamp import TIME_GRANULARITY
  
  # AfterCount is experimental. No backwards compatibility guarantees.
  
  __all__ = [
      'AccumulationMode',
      'TriggerFn',
      'DefaultTrigger',
      'AfterWatermark',
      'AfterProcessingTime',
      'AfterCount',
      'Repeatedly',
      'AfterAny',
      'AfterAll',
      'AfterEach',
      'OrFinally',
  ]
  
  _LOGGER = logging.getLogger(__name__)
  
  
  class AccumulationMode(object):
    """Controls what to do with data when a trigger fires multiple times."""
    DISCARDING = beam_runner_api_pb2.AccumulationMode.DISCARDING
    ACCUMULATING = beam_runner_api_pb2.AccumulationMode.ACCUMULATING
    # TODO(robertwb): Provide retractions of previous outputs.
    # RETRACTING = 3
  
  
  class _StateTag(with_metaclass(ABCMeta, object)):  # type: ignore[misc]
    """An identifier used to store and retrieve typed, combinable state.
  
    The given tag must be unique for this step."""
    def __init__(self, tag):
      self.tag = tag
  
  
  class _ValueStateTag(_StateTag):
    """StateTag pointing to an element."""
    def __repr__(self):
      return 'ValueStateTag(%s)' % (self.tag)
  
    def with_prefix(self, prefix):
      return _ValueStateTag(prefix + self.tag)
  
  
  class _SetStateTag(_StateTag):
    """StateTag pointing to an element."""
    def __repr__(self):
      return 'SetStateTag({tag})'.format(tag=self.tag)
  
    def with_prefix(self, prefix):
      return _SetStateTag(prefix + self.tag)
  
  
  class _CombiningValueStateTag(_StateTag):
    """StateTag pointing to an element, accumulated with a combiner.
  
    The given tag must be unique for this step. The given CombineFn will be
    applied (possibly incrementally and eagerly) when adding elements."""
  
    # TODO(robertwb): Also store the coder (perhaps extracted from the combine_fn)
    def __init__(self, tag, combine_fn):
      super(_CombiningValueStateTag, self).__init__(tag)
      if not combine_fn:
        raise ValueError('combine_fn must be specified.')
      if not isinstance(combine_fn, core.CombineFn):
        combine_fn = core.CombineFn.from_callable(combine_fn)
      self.combine_fn = combine_fn
  
    def __repr__(self):
      return 'CombiningValueStateTag(%s, %s)' % (self.tag, self.combine_fn)
  
    def with_prefix(self, prefix):
      return _CombiningValueStateTag(prefix + self.tag, self.combine_fn)
  
    def without_extraction(self):
      class NoExtractionCombineFn(core.CombineFn):
        create_accumulator = self.combine_fn.create_accumulator
        add_input = self.combine_fn.add_input
        merge_accumulators = self.combine_fn.merge_accumulators
        compact = self.combine_fn.compact
        extract_output = staticmethod(lambda x: x)
  
      return _CombiningValueStateTag(self.tag, NoExtractionCombineFn())
  
  
  class _ListStateTag(_StateTag):
    """StateTag pointing to a list of elements."""
    def __repr__(self):
      return 'ListStateTag(%s)' % self.tag
  
    def with_prefix(self, prefix):
      return _ListStateTag(prefix + self.tag)
  
  
  class _WatermarkHoldStateTag(_StateTag):
    def __init__(self, tag, timestamp_combiner_impl):
      super(_WatermarkHoldStateTag, self).__init__(tag)
      self.timestamp_combiner_impl = timestamp_combiner_impl
  
    def __repr__(self):
      return 'WatermarkHoldStateTag(%s, %s)' % (
          self.tag, self.timestamp_combiner_impl)
  
    def with_prefix(self, prefix):
      return _WatermarkHoldStateTag(
          prefix + self.tag, self.timestamp_combiner_impl)
  
  
  # pylint: disable=unused-argument
  # TODO(robertwb): Provisional API, Java likely to change as well.
  class TriggerFn(with_metaclass(ABCMeta, object)):  # type: ignore[misc]
    """A TriggerFn determines when window (panes) are emitted.
  
    See https://beam.apache.org/documentation/programming-guide/#triggers
    """
    @abstractmethod
    def on_element(self, element, window, context):
      """Called when a new element arrives in a window.
  
      Args:
        element: the element being added
        window: the window to which the element is being added
        context: a context (e.g. a TriggerContext instance) for managing state
            and setting timers
      """
      pass
  
    @abstractmethod
    def on_merge(self, to_be_merged, merge_result, context):
      """Called when multiple windows are merged.
  
      Args:
        to_be_merged: the set of windows to be merged
        merge_result: the window into which the windows are being merged
        context: a context (e.g. a TriggerContext instance) for managing state
            and setting timers
      """
      pass
  
    @abstractmethod
    def should_fire(self, time_domain, timestamp, window, context):
      """Whether this trigger should cause the window to fire.
  
      Args:
        time_domain: WATERMARK for event-time timers and REAL_TIME for
            processing-time timers.
        timestamp: for time_domain WATERMARK, it represents the
            watermark: (a lower bound on) the watermark of the system
            and for time_domain REAL_TIME, it represents the
            trigger: timestamp of the processing-time timer.
        window: the window whose trigger is being considered
        context: a context (e.g. a TriggerContext instance) for managing state
            and setting timers
  
      Returns:
        whether this trigger should cause a firing
      """
      pass
  
    @abstractmethod
    def has_ontime_pane(self):
      """Whether this trigger creates an empty pane even if there are no elements.
  
      Returns:
        True if this trigger guarantees that there will always be an ON_TIME pane
        even if there are no elements in that pane.
      """
      pass
  
    @abstractmethod
    def on_fire(self, watermark, window, context):
      """Called when a trigger actually fires.
  
      Args:
        watermark: (a lower bound on) the watermark of the system
        window: the window whose trigger is being fired
        context: a context (e.g. a TriggerContext instance) for managing state
            and setting timers
  
      Returns:
        whether this trigger is finished
      """
      pass
  
    @abstractmethod
    def reset(self, window, context):
      """Clear any state and timers used by this TriggerFn."""
      pass
  
  
  # pylint: enable=unused-argument
  
    @staticmethod
    def from_runner_api(proto, context):
      return {
          'after_all': AfterAll,
          'after_any': AfterAny,
          'after_each': AfterEach,
          'after_end_of_window': AfterWatermark,
          'after_processing_time': AfterProcessingTime,
          # after_processing_time, after_synchronized_processing_time
          'always': Always,
          'default': DefaultTrigger,
          'element_count': AfterCount,
          # never
          'or_finally': OrFinally,
          'repeat': Repeatedly,
      }[proto.WhichOneof('trigger')].from_runner_api(proto, context)
  
    @abstractmethod
    def to_runner_api(self, unused_context):
      pass
  
    def __ne__(self, other):
      # TODO(BEAM-5949): Needed for Python 2 compatibility.
      return not self == other
  
  
  class DefaultTrigger(TriggerFn):
    """Semantically Repeatedly(AfterWatermark()), but more optimized."""
    def __init__(self):
      pass
  
    def __repr__(self):
      return 'DefaultTrigger()'
  
    def on_element(self, element, window, context):
      context.set_timer('', TimeDomain.WATERMARK, window.end)
  
    def on_merge(self, to_be_merged, merge_result, context):
      # Note: Timer clearing solely an optimization.
      for window in to_be_merged:
        if window.end != merge_result.end:
          context.clear_timer('', TimeDomain.WATERMARK)
  
    def should_fire(self, time_domain, watermark, window, context):
      if watermark >= window.end:
        # Explicitly clear the timer so that late elements are not emitted again
        # when the timer is fired.
        context.clear_timer('', TimeDomain.WATERMARK)
      return watermark >= window.end
  
    def on_fire(self, watermark, window, context):
      return False
  
    def reset(self, window, context):
      context.clear_timer('', TimeDomain.WATERMARK)
  
    def __eq__(self, other):
      return type(self) == type(other)
  
    def __hash__(self):
      return hash(type(self))
  
    @staticmethod
    def from_runner_api(proto, context):
      return DefaultTrigger()
  
    def to_runner_api(self, unused_context):
      return beam_runner_api_pb2.Trigger(
          default=beam_runner_api_pb2.Trigger.Default())
  
    def has_ontime_pane(self):
      return True
  
  
  class AfterProcessingTime(TriggerFn):
    """Fire exactly once after a specified delay from processing time.
  
    AfterProcessingTime is experimental. No backwards compatibility guarantees.
    """
    def __init__(self, delay=0):
      """Initialize a processing time trigger with a delay in seconds."""
      self.delay = delay
  
    def __repr__(self):
      return 'AfterProcessingTime(delay=%d)' % self.delay
  
    def on_element(self, element, window, context):
      context.set_timer(
          '', TimeDomain.REAL_TIME, context.get_current_time() + self.delay)
  
    def on_merge(self, to_be_merged, merge_result, context):
      # timers will be kept through merging
      pass
  
    def should_fire(self, time_domain, timestamp, window, context):
      if time_domain == TimeDomain.REAL_TIME:
        return True
  
    def on_fire(self, timestamp, window, context):
      return True
  
    def reset(self, window, context):
      pass
  
    @staticmethod
    def from_runner_api(proto, context):
      return AfterProcessingTime(
          delay=(
              proto.after_processing_time.timestamp_transforms[0].delay.
              delay_millis) // 1000)
  
    def to_runner_api(self, context):
      delay_proto = beam_runner_api_pb2.TimestampTransform(
          delay=beam_runner_api_pb2.TimestampTransform.Delay(
              delay_millis=self.delay * 1000))
      return beam_runner_api_pb2.Trigger(
          after_processing_time=beam_runner_api_pb2.Trigger.AfterProcessingTime(
              timestamp_transforms=[delay_proto]))
  
    def has_ontime_pane(self):
      return False
  
  
  class Always(TriggerFn):
    """Repeatedly invoke the given trigger, never finishing."""
    def __init__(self):
      pass
  
    def __repr__(self):
      return 'Always'
  
    def __eq__(self, other):
      return type(self) == type(other)
  
    def __hash__(self):
      return 1
  
    def on_element(self, element, window, context):
      pass
  
    def on_merge(self, to_be_merged, merge_result, context):
      pass
  
    def has_ontime_pane(self):
      return False
  
    def reset(self, window, context):
      pass
  
    def should_fire(self, time_domain, watermark, window, context):
      return True
  
    def on_fire(self, watermark, window, context):
      return False
  
    @staticmethod
    def from_runner_api(proto, context):
      return Always()
  
    def to_runner_api(self, context):
      return beam_runner_api_pb2.Trigger(
          always=beam_runner_api_pb2.Trigger.Always())
  
  
  class AfterWatermark(TriggerFn):
    """Fire exactly once when the watermark passes the end of the window.
  
    Args:
        early: if not None, a speculative trigger to repeatedly evaluate before
          the watermark passes the end of the window
        late: if not None, a speculative trigger to repeatedly evaluate after
          the watermark passes the end of the window
    """
    LATE_TAG = _CombiningValueStateTag('is_late', any)
  
    def __init__(self, early=None, late=None):
      self.early = Repeatedly(early) if early else None
      self.late = Repeatedly(late) if late else None
  
    def __repr__(self):
      qualifiers = []
      if self.early:
        qualifiers.append('early=%s' % self.early.underlying)
      if self.late:
        qualifiers.append('late=%s' % self.late.underlying)
      return 'AfterWatermark(%s)' % ', '.join(qualifiers)
  
    def is_late(self, context):
      return self.late and context.get_state(self.LATE_TAG)
  
    def on_element(self, element, window, context):
      if self.is_late(context):
        self.late.on_element(element, window, NestedContext(context, 'late'))
      else:
        context.set_timer('', TimeDomain.WATERMARK, window.end)
        if self.early:
          self.early.on_element(element, window, NestedContext(context, 'early'))
  
    def on_merge(self, to_be_merged, merge_result, context):
      # TODO(robertwb): Figure out whether the 'rewind' semantics could be used
      # here.
      if self.is_late(context):
        self.late.on_merge(
            to_be_merged, merge_result, NestedContext(context, 'late'))
      else:
        # Note: Timer clearing solely an optimization.
        for window in to_be_merged:
          if window.end != merge_result.end:
            context.clear_timer('', TimeDomain.WATERMARK)
        if self.early:
          self.early.on_merge(
              to_be_merged, merge_result, NestedContext(context, 'early'))
  
    def should_fire(self, time_domain, watermark, window, context):
      if self.is_late(context):
        return self.late.should_fire(
            time_domain, watermark, window, NestedContext(context, 'late'))
      elif watermark >= window.end:
        # Explicitly clear the timer so that late elements are not emitted again
        # when the timer is fired.
        context.clear_timer('', TimeDomain.WATERMARK)
        return True
      elif self.early:
        return self.early.should_fire(
            time_domain, watermark, window, NestedContext(context, 'early'))
      return False
  
    def on_fire(self, watermark, window, context):
      if self.is_late(context):
        return self.late.on_fire(
            watermark, window, NestedContext(context, 'late'))
      elif watermark >= window.end:
        context.add_state(self.LATE_TAG, True)
        return not self.late
      elif self.early:
        self.early.on_fire(watermark, window, NestedContext(context, 'early'))
        return False
  
    def reset(self, window, context):
      if self.late:
        context.clear_state(self.LATE_TAG)
      if self.early:
        self.early.reset(window, NestedContext(context, 'early'))
      if self.late:
        self.late.reset(window, NestedContext(context, 'late'))
  
    def __eq__(self, other):
      return (
          type(self) == type(other) and self.early == other.early and
          self.late == other.late)
  
    def __hash__(self):
      return hash((type(self), self.early, self.late))
  
    @staticmethod
    def from_runner_api(proto, context):
      return AfterWatermark(
          early=TriggerFn.from_runner_api(
              proto.after_end_of_window.early_firings, context)
          if proto.after_end_of_window.HasField('early_firings') else None,
          late=TriggerFn.from_runner_api(
              proto.after_end_of_window.late_firings, context)
          if proto.after_end_of_window.HasField('late_firings') else None)
  
    def to_runner_api(self, context):
      early_proto = self.early.underlying.to_runner_api(
          context) if self.early else None
      late_proto = self.late.underlying.to_runner_api(
          context) if self.late else None
      return beam_runner_api_pb2.Trigger(
          after_end_of_window=beam_runner_api_pb2.Trigger.AfterEndOfWindow(
              early_firings=early_proto, late_firings=late_proto))
  
    def has_ontime_pane(self):
      return True
  
  
  class AfterCount(TriggerFn):
    """Fire when there are at least count elements in this window pane.
  
    AfterCount is experimental. No backwards compatibility guarantees.
    """
  
    COUNT_TAG = _CombiningValueStateTag('count', combiners.CountCombineFn())
  
    def __init__(self, count):
      if not isinstance(count, numbers.Integral) or count < 1:
        raise ValueError("count (%d) must be a positive integer." % count)
      self.count = count
  
    def __repr__(self):
      return 'AfterCount(%s)' % self.count
  
    def __eq__(self, other):
      return type(self) == type(other) and self.count == other.count
  
    def __hash__(self):
      return hash(self.count)
  
    def on_element(self, element, window, context):
      context.add_state(self.COUNT_TAG, 1)
  
    def on_merge(self, to_be_merged, merge_result, context):
      # states automatically merged
      pass
  
    def should_fire(self, time_domain, watermark, window, context):
      return context.get_state(self.COUNT_TAG) >= self.count
  
    def on_fire(self, watermark, window, context):
      return True
  
    def reset(self, window, context):
      context.clear_state(self.COUNT_TAG)
  
    @staticmethod
    def from_runner_api(proto, unused_context):
      return AfterCount(proto.element_count.element_count)
  
    def to_runner_api(self, unused_context):
      return beam_runner_api_pb2.Trigger(
          element_count=beam_runner_api_pb2.Trigger.ElementCount(
              element_count=self.count))
  
    def has_ontime_pane(self):
      return False
  
  
  class Repeatedly(TriggerFn):
    """Repeatedly invoke the given trigger, never finishing."""
    def __init__(self, underlying):
      self.underlying = underlying
  
    def __repr__(self):
      return 'Repeatedly(%s)' % self.underlying
  
    def __eq__(self, other):
      return type(self) == type(other) and self.underlying == other.underlying
  
    def __hash__(self):
      return hash(self.underlying)
  
    def on_element(self, element, window, context):
      self.underlying.on_element(element, window, context)
  
    def on_merge(self, to_be_merged, merge_result, context):
      self.underlying.on_merge(to_be_merged, merge_result, context)
  
    def should_fire(self, time_domain, watermark, window, context):
      return self.underlying.should_fire(time_domain, watermark, window, context)
  
    def on_fire(self, watermark, window, context):
      if self.underlying.on_fire(watermark, window, context):
        self.underlying.reset(window, context)
      return False
  
    def reset(self, window, context):
      self.underlying.reset(window, context)
  
    @staticmethod
    def from_runner_api(proto, context):
      return Repeatedly(
          TriggerFn.from_runner_api(proto.repeat.subtrigger, context))
  
    def to_runner_api(self, context):
      return beam_runner_api_pb2.Trigger(
          repeat=beam_runner_api_pb2.Trigger.Repeat(
              subtrigger=self.underlying.to_runner_api(context)))
  
    def has_ontime_pane(self):
      return self.underlying.has_ontime_pane()
  
  
  class _ParallelTriggerFn(with_metaclass(ABCMeta,
                                          TriggerFn)):  # type: ignore[misc]
    def __init__(self, *triggers):
      self.triggers = triggers
  
    def __repr__(self):
      return '%s(%s)' % (
          self.__class__.__name__, ', '.join(str(t) for t in self.triggers))
  
    def __eq__(self, other):
      return type(self) == type(other) and self.triggers == other.triggers
  
    def __hash__(self):
      return hash(self.triggers)
  
    @abstractmethod
    def combine_op(self, trigger_results):
      pass
  
    def on_element(self, element, window, context):
      for ix, trigger in enumerate(self.triggers):
        trigger.on_element(element, window, self._sub_context(context, ix))
  
    def on_merge(self, to_be_merged, merge_result, context):
      for ix, trigger in enumerate(self.triggers):
        trigger.on_merge(
            to_be_merged, merge_result, self._sub_context(context, ix))
  
    def should_fire(self, time_domain, watermark, window, context):
      self._time_domain = time_domain
      return self.combine_op(
          trigger.should_fire(
              time_domain, watermark, window, self._sub_context(context, ix))
          for ix,
          trigger in enumerate(self.triggers))
  
    def on_fire(self, watermark, window, context):
      finished = []
      for ix, trigger in enumerate(self.triggers):
        nested_context = self._sub_context(context, ix)
        if trigger.should_fire(TimeDomain.WATERMARK,
                               watermark,
                               window,
                               nested_context):
          finished.append(trigger.on_fire(watermark, window, nested_context))
      return self.combine_op(finished)
  
    def reset(self, window, context):
      for ix, trigger in enumerate(self.triggers):
        trigger.reset(window, self._sub_context(context, ix))
  
    @staticmethod
    def _sub_context(context, index):
      return NestedContext(context, '%d/' % index)
  
    @staticmethod
    def from_runner_api(proto, context):
      subtriggers = [
          TriggerFn.from_runner_api(subtrigger, context) for subtrigger in
          proto.after_all.subtriggers or proto.after_any.subtriggers
      ]
      if proto.after_all.subtriggers:
        return AfterAll(*subtriggers)
      else:
        return AfterAny(*subtriggers)
  
    def to_runner_api(self, context):
      subtriggers = [
          subtrigger.to_runner_api(context) for subtrigger in self.triggers
      ]
      if self.combine_op == all:
        return beam_runner_api_pb2.Trigger(
            after_all=beam_runner_api_pb2.Trigger.AfterAll(
                subtriggers=subtriggers))
      elif self.combine_op == any:
        return beam_runner_api_pb2.Trigger(
            after_any=beam_runner_api_pb2.Trigger.AfterAny(
                subtriggers=subtriggers))
      else:
        raise NotImplementedError(self)
  
    def has_ontime_pane(self):
      return any(t.has_ontime_pane() for t in self.triggers)
  
  
  class AfterAny(_ParallelTriggerFn):
    """Fires when any subtrigger fires.
  
    Also finishes when any subtrigger finishes.
    """
    combine_op = any
  
  
  class AfterAll(_ParallelTriggerFn):
    """Fires when all subtriggers have fired.
  
    Also finishes when all subtriggers have finished.
    """
    combine_op = all
  
  
  class AfterEach(TriggerFn):
  
    INDEX_TAG = _CombiningValueStateTag(
        'index', (lambda indices: 0 if not indices else max(indices)))
  
    def __init__(self, *triggers):
      self.triggers = triggers
  
    def __repr__(self):
      return '%s(%s)' % (
          self.__class__.__name__, ', '.join(str(t) for t in self.triggers))
  
    def __eq__(self, other):
      return type(self) == type(other) and self.triggers == other.triggers
  
    def __hash__(self):
      return hash(self.triggers)
  
    def on_element(self, element, window, context):
      ix = context.get_state(self.INDEX_TAG)
      if ix < len(self.triggers):
        self.triggers[ix].on_element(
            element, window, self._sub_context(context, ix))
  
    def on_merge(self, to_be_merged, merge_result, context):
      # This takes the furthest window on merging.
      # TODO(robertwb): Revisit this when merging windows logic is settled for
      # all possible merging situations.
      ix = context.get_state(self.INDEX_TAG)
      if ix < len(self.triggers):
        self.triggers[ix].on_merge(
            to_be_merged, merge_result, self._sub_context(context, ix))
  
    def should_fire(self, time_domain, watermark, window, context):
      ix = context.get_state(self.INDEX_TAG)
      if ix < len(self.triggers):
        return self.triggers[ix].should_fire(
            time_domain, watermark, window, self._sub_context(context, ix))
  
    def on_fire(self, watermark, window, context):
      ix = context.get_state(self.INDEX_TAG)
      if ix < len(self.triggers):
        if self.triggers[ix].on_fire(watermark,
                                     window,
                                     self._sub_context(context, ix)):
          ix += 1
          context.add_state(self.INDEX_TAG, ix)
        return ix == len(self.triggers)
  
    def reset(self, window, context):
      context.clear_state(self.INDEX_TAG)
      for ix, trigger in enumerate(self.triggers):
        trigger.reset(window, self._sub_context(context, ix))
  
    @staticmethod
    def _sub_context(context, index):
      return NestedContext(context, '%d/' % index)
  
    @staticmethod
    def from_runner_api(proto, context):
      return AfterEach(
          *[
              TriggerFn.from_runner_api(subtrigger, context)
              for subtrigger in proto.after_each.subtriggers
          ])
  
    def to_runner_api(self, context):
      return beam_runner_api_pb2.Trigger(
          after_each=beam_runner_api_pb2.Trigger.AfterEach(
              subtriggers=[
                  subtrigger.to_runner_api(context)
                  for subtrigger in self.triggers
              ]))
  
    def has_ontime_pane(self):
      return any(t.has_ontime_pane() for t in self.triggers)
  
  
  class OrFinally(AfterAny):
    @staticmethod
    def from_runner_api(proto, context):
      return OrFinally(
          TriggerFn.from_runner_api(proto.or_finally.main, context),
          # getattr is used as finally is a keyword in Python
          TriggerFn.from_runner_api(
              getattr(proto.or_finally, 'finally'), context))
  
    def to_runner_api(self, context):
      return beam_runner_api_pb2.Trigger(
          or_finally=beam_runner_api_pb2.Trigger.OrFinally(
              main=self.triggers[0].to_runner_api(context),
              # dict keyword argument is used as finally is a keyword in Python
              **{'finally': self.triggers[1].to_runner_api(context)}))
  
  
  class TriggerContext(object):
    def __init__(self, outer, window, clock):
      self._outer = outer
      self._window = window
      self._clock = clock
  
    def get_current_time(self):
      return self._clock.time()
  
    def set_timer(self, name, time_domain, timestamp):
      self._outer.set_timer(self._window, name, time_domain, timestamp)
  
    def clear_timer(self, name, time_domain):
      self._outer.clear_timer(self._window, name, time_domain)
  
    def add_state(self, tag, value):
      self._outer.add_state(self._window, tag, value)
  
    def get_state(self, tag):
      return self._outer.get_state(self._window, tag)
  
    def clear_state(self, tag):
      return self._outer.clear_state(self._window, tag)
  
  
  class NestedContext(object):
    """Namespaced context useful for defining composite triggers."""
    def __init__(self, outer, prefix):
      self._outer = outer
      self._prefix = prefix
  
    def get_current_time(self):
      return self._outer.get_current_time()
  
    def set_timer(self, name, time_domain, timestamp):
      self._outer.set_timer(self._prefix + name, time_domain, timestamp)
  
    def clear_timer(self, name, time_domain):
      self._outer.clear_timer(self._prefix + name, time_domain)
  
    def add_state(self, tag, value):
      self._outer.add_state(tag.with_prefix(self._prefix), value)
  
    def get_state(self, tag):
      return self._outer.get_state(tag.with_prefix(self._prefix))
  
    def clear_state(self, tag):
      self._outer.clear_state(tag.with_prefix(self._prefix))
  
  
  # pylint: disable=unused-argument
  class SimpleState(with_metaclass(ABCMeta, object)):  # type: ignore[misc]
    """Basic state storage interface used for triggering.
  
    Only timers must hold the watermark (by their timestamp).
    """
    @abstractmethod
    def set_timer(self, window, name, time_domain, timestamp):
      pass
  
    @abstractmethod
    def get_window(self, window_id):
      pass
  
    @abstractmethod
    def clear_timer(self, window, name, time_domain):
      pass
  
    @abstractmethod
    def add_state(self, window, tag, value):
      pass
  
    @abstractmethod
    def get_state(self, window, tag):
      pass
  
    @abstractmethod
    def clear_state(self, window, tag):
      pass
  
    def at(self, window, clock):
      return NestedContext(TriggerContext(self, window, clock), 'trigger')
  
  
  class UnmergedState(SimpleState):
    """State suitable for use in TriggerDriver.
  
    This class must be implemented by each backend.
    """
    @abstractmethod
    def set_global_state(self, tag, value):
      pass
  
    @abstractmethod
    def get_global_state(self, tag, default=None):
      pass
  
  
  # pylint: enable=unused-argument
  
  
  class MergeableStateAdapter(SimpleState):
    """Wraps an UnmergedState, tracking merged windows."""
    # TODO(robertwb): A similar indirection could be used for sliding windows
    # or other window_fns when a single element typically belongs to many windows.
  
    WINDOW_IDS = _ValueStateTag('window_ids')
  
    def __init__(self, raw_state):
      self.raw_state = raw_state
      self.window_ids = self.raw_state.get_global_state(self.WINDOW_IDS, {})
      self.counter = None
  
    def set_timer(self, window, name, time_domain, timestamp):
      self.raw_state.set_timer(self._get_id(window), name, time_domain, timestamp)
  
    def clear_timer(self, window, name, time_domain):
      for window_id in self._get_ids(window):
        self.raw_state.clear_timer(window_id, name, time_domain)
  
    def add_state(self, window, tag, value):
      if isinstance(tag, _ValueStateTag):
        raise ValueError(
            'Merging requested for non-mergeable state tag: %r.' % tag)
      elif isinstance(tag, _CombiningValueStateTag):
        tag = tag.without_extraction()
      self.raw_state.add_state(self._get_id(window), tag, value)
  
    def get_state(self, window, tag):
      if isinstance(tag, _CombiningValueStateTag):
        original_tag, tag = tag, tag.without_extraction()
      values = [
          self.raw_state.get_state(window_id, tag)
          for window_id in self._get_ids(window)
      ]
      if isinstance(tag, _ValueStateTag):
        raise ValueError(
            'Merging requested for non-mergeable state tag: %r.' % tag)
      elif isinstance(tag, _CombiningValueStateTag):
        return original_tag.combine_fn.extract_output(
            original_tag.combine_fn.merge_accumulators(values))
      elif isinstance(tag, _ListStateTag):
        return [v for vs in values for v in vs]
      elif isinstance(tag, _SetStateTag):
        return {v for vs in values for v in vs}
      elif isinstance(tag, _WatermarkHoldStateTag):
        return tag.timestamp_combiner_impl.combine_all(values)
      else:
        raise ValueError('Invalid tag.', tag)
  
    def clear_state(self, window, tag):
      for window_id in self._get_ids(window):
        self.raw_state.clear_state(window_id, tag)
      if tag is None:
        del self.window_ids[window]
        self._persist_window_ids()
  
    def merge(self, to_be_merged, merge_result):
      for window in to_be_merged:
        if window != merge_result:
          if window in self.window_ids:
            if merge_result in self.window_ids:
              merge_window_ids = self.window_ids[merge_result]
            else:
              merge_window_ids = self.window_ids[merge_result] = []
            merge_window_ids.extend(self.window_ids.pop(window))
            self._persist_window_ids()
  
    def known_windows(self):
      return list(self.window_ids)
  
    def get_window(self, window_id):
      for window, ids in self.window_ids.items():
        if window_id in ids:
          return window
      raise ValueError('No window for %s' % window_id)
  
    def _get_id(self, window):
      if window in self.window_ids:
        return self.window_ids[window][0]
  
      window_id = self._get_next_counter()
      self.window_ids[window] = [window_id]
      self._persist_window_ids()
      return window_id
  
    def _get_ids(self, window):
      return self.window_ids.get(window, [])
  
    def _get_next_counter(self):
      if not self.window_ids:
        self.counter = 0
      elif self.counter is None:
        self.counter = max(k for ids in self.window_ids.values() for k in ids)
      self.counter += 1
      return self.counter
  
    def _persist_window_ids(self):
      self.raw_state.set_global_state(self.WINDOW_IDS, self.window_ids)
  
    def __repr__(self):
      return '\n\t'.join([repr(self.window_ids)] +
                         repr(self.raw_state).split('\n'))
  
  
  def create_trigger_driver(
      windowing, is_batch=False, phased_combine_fn=None, clock=None):
    """Create the TriggerDriver for the given windowing and options."""
  
    # TODO(robertwb): We can do more if we know elements are in timestamp
    # sorted order.
    if windowing.is_default() and is_batch:
      driver = BatchGlobalTriggerDriver()
    elif (windowing.windowfn == GlobalWindows() and
          (windowing.triggerfn in [AfterCount(1), Always()]) and is_batch):
      # Here we also just pass through all the values exactly once.
      driver = BatchGlobalTriggerDriver()
    else:
      driver = GeneralTriggerDriver(windowing, clock)
  
    if phased_combine_fn:
      # TODO(ccy): Refactor GeneralTriggerDriver to combine values eagerly using
      # the known phased_combine_fn here.
      driver = CombiningTriggerDriver(phased_combine_fn, driver)
    return driver
  
  
  class TriggerDriver(with_metaclass(ABCMeta, object)):  # type: ignore[misc]
    """Breaks a series of bundle and timer firings into window (pane)s."""
    @abstractmethod
    def process_elements(
        self,
        state,
        windowed_values,
        output_watermark,
        input_watermark=MIN_TIMESTAMP):
      pass
  
    @abstractmethod
    def process_timer(
        self,
        window_id,
        name,
        time_domain,
        timestamp,
        state,
        input_watermark=None):
      pass
  
    def process_entire_key(self, key, windowed_values):
      state = InMemoryUnmergedState()
      for wvalue in self.process_elements(state,
                                          windowed_values,
                                          MIN_TIMESTAMP,
                                          MIN_TIMESTAMP):
        yield wvalue.with_value((key, wvalue.value))
      while state.timers:
        fired = state.get_and_clear_timers()
        for timer_window, (name, time_domain, fire_time) in fired:
          for wvalue in self.process_timer(timer_window,
                                           name,
                                           time_domain,
                                           fire_time,
                                           state):
            yield wvalue.with_value((key, wvalue.value))
  
  
  class _UnwindowedValues(observable.ObservableMixin):
    """Exposes iterable of windowed values as iterable of unwindowed values."""
    def __init__(self, windowed_values):
      super(_UnwindowedValues, self).__init__()
      self._windowed_values = windowed_values
  
    def __iter__(self):
      for wv in self._windowed_values:
        unwindowed_value = wv.value
        self.notify_observers(unwindowed_value)
        yield unwindowed_value
  
    def __repr__(self):
      return '<_UnwindowedValues of %s>' % self._windowed_values
  
    def __reduce__(self):
      return list, (list(self), )
  
    def __eq__(self, other):
      if isinstance(other, collections.Iterable):
        return all(
            a == b for a, b in zip_longest(self, other, fillvalue=object()))
      else:
        return NotImplemented
  
    def __hash__(self):
      return hash(tuple(self))
  
    def __ne__(self, other):
      # TODO(BEAM-5949): Needed for Python 2 compatibility.
      return not self == other
  
  
  coder_impl.FastPrimitivesCoderImpl.register_iterable_like_type(
      _UnwindowedValues)
  
  
  class BatchGlobalTriggerDriver(TriggerDriver):
    """Groups all received values together.
    """
    GLOBAL_WINDOW_TUPLE = (GlobalWindow(), )
    ONLY_FIRING = windowed_value.PaneInfo(
        is_first=True,
        is_last=True,
        timing=windowed_value.PaneInfoTiming.ON_TIME,
        index=0,
        nonspeculative_index=0)
  
    def process_elements(
        self,
        state,
        windowed_values,
        unused_output_watermark,
        unused_input_watermark=MIN_TIMESTAMP):
      yield WindowedValue(
          _UnwindowedValues(windowed_values),
          MIN_TIMESTAMP,
          self.GLOBAL_WINDOW_TUPLE,
          self.ONLY_FIRING)
  
    def process_timer(
        self,
        window_id,
        name,
        time_domain,
        timestamp,
        state,
        input_watermark=None):
      raise TypeError('Triggers never set or called for batch default windowing.')
  
  
  class CombiningTriggerDriver(TriggerDriver):
    """Uses a phased_combine_fn to process output of wrapped TriggerDriver."""
    def __init__(self, phased_combine_fn, underlying):
      self.phased_combine_fn = phased_combine_fn
      self.underlying = underlying
  
    def process_elements(
        self,
        state,
        windowed_values,
        output_watermark,
        input_watermark=MIN_TIMESTAMP):
      uncombined = self.underlying.process_elements(
          state, windowed_values, output_watermark, input_watermark)
      for output in uncombined:
        yield output.with_value(self.phased_combine_fn.apply(output.value))
  
    def process_timer(
        self,
        window_id,
        name,
        time_domain,
        timestamp,
        state,
        input_watermark=None):
      uncombined = self.underlying.process_timer(
          window_id, name, time_domain, timestamp, state, input_watermark)
      for output in uncombined:
        yield output.with_value(self.phased_combine_fn.apply(output.value))
  
  
  class GeneralTriggerDriver(TriggerDriver):
    """Breaks a series of bundle and timer firings into window (pane)s.
  
    Suitable for all variants of Windowing.
    """
    ELEMENTS = _ListStateTag('elements')
    TOMBSTONE = _CombiningValueStateTag('tombstone', combiners.CountCombineFn())
    INDEX = _CombiningValueStateTag('index', combiners.CountCombineFn())
    NONSPECULATIVE_INDEX = _CombiningValueStateTag(
        'nonspeculative_index', combiners.CountCombineFn())
  
    def __init__(self, windowing, clock):
      self.clock = clock
      self.allowed_lateness = windowing.allowed_lateness
      self.window_fn = windowing.windowfn
      self.timestamp_combiner_impl = TimestampCombiner.get_impl(
          windowing.timestamp_combiner, self.window_fn)
      # pylint: disable=invalid-name
      self.WATERMARK_HOLD = _WatermarkHoldStateTag(
          'watermark', self.timestamp_combiner_impl)
      # pylint: enable=invalid-name
      self.trigger_fn = windowing.triggerfn
      self.accumulation_mode = windowing.accumulation_mode
      self.is_merging = True
  
    def process_elements(
        self,
        state,
        windowed_values,
        output_watermark,
        input_watermark=MIN_TIMESTAMP):
      if self.is_merging:
        state = MergeableStateAdapter(state)
  
      windows_to_elements = collections.defaultdict(list)
      for wv in windowed_values:
        for window in wv.windows:
          # ignore expired windows
          if input_watermark > window.end + self.allowed_lateness:
            continue
          windows_to_elements[window].append((wv.value, wv.timestamp))
  
      # First handle merging.
      if self.is_merging:
        old_windows = set(state.known_windows())
        all_windows = old_windows.union(list(windows_to_elements))
  
        if all_windows != old_windows:
          merged_away = {}
  
          class TriggerMergeContext(WindowFn.MergeContext):
            def merge(_, to_be_merged, merge_result):  # pylint: disable=no-self-argument
              for window in to_be_merged:
                if window != merge_result:
                  merged_away[window] = merge_result
                  # Clear state associated with PaneInfo since it is
                  # not preserved across merges.
                  state.clear_state(window, self.INDEX)
                  state.clear_state(window, self.NONSPECULATIVE_INDEX)
              state.merge(to_be_merged, merge_result)
              # using the outer self argument.
              self.trigger_fn.on_merge(
                  to_be_merged, merge_result, state.at(merge_result, self.clock))
  
          self.window_fn.merge(TriggerMergeContext(all_windows))
  
          merged_windows_to_elements = collections.defaultdict(list)
          for window, values in windows_to_elements.items():
            while window in merged_away:
              window = merged_away[window]
            merged_windows_to_elements[window].extend(values)
          windows_to_elements = merged_windows_to_elements
  
          for window in merged_away:
            state.clear_state(window, self.WATERMARK_HOLD)
  
      # Next handle element adding.
      for window, elements in windows_to_elements.items():
        if state.get_state(window, self.TOMBSTONE):
          continue
        # Add watermark hold.
        # TODO(ccy): Add late data and garbage-collection hold support.
        output_time = self.timestamp_combiner_impl.merge(
            window,
            (
                element_output_time for element_output_time in (
                    self.timestamp_combiner_impl.assign_output_time(
                        window, timestamp) for unused_value,
                    timestamp in elements)
                if element_output_time >= output_watermark))
        if output_time is not None:
          state.add_state(window, self.WATERMARK_HOLD, output_time)
  
        context = state.at(window, self.clock)
        for value, unused_timestamp in elements:
          state.add_state(window, self.ELEMENTS, value)
          self.trigger_fn.on_element(value, window, context)
  
        # Maybe fire this window.
        if self.trigger_fn.should_fire(TimeDomain.WATERMARK,
                                       input_watermark,
                                       window,
                                       context):
          finished = self.trigger_fn.on_fire(input_watermark, window, context)
          yield self._output(window, finished, state, output_watermark, False)
  
    def process_timer(
        self,
        window_id,
        unused_name,
        time_domain,
        timestamp,
        state,
        input_watermark=None):
      if input_watermark is None:
        input_watermark = timestamp
  
      if self.is_merging:
        state = MergeableStateAdapter(state)
      window = state.get_window(window_id)
      if state.get_state(window, self.TOMBSTONE):
        return
  
      if time_domain in (TimeDomain.WATERMARK, TimeDomain.REAL_TIME):
        if not self.is_merging or window in state.known_windows():
          context = state.at(window, self.clock)
          if self.trigger_fn.should_fire(time_domain, timestamp, window, context):
            finished = self.trigger_fn.on_fire(timestamp, window, context)
            yield self._output(
                window,
                finished,
                state,
                timestamp,
                time_domain == TimeDomain.WATERMARK)
      else:
        raise Exception('Unexpected time domain: %s' % time_domain)
  
    def _output(self, window, finished, state, output_watermark, maybe_ontime):
      """Output window and clean up if appropriate."""
      index = state.get_state(window, self.INDEX)
      state.add_state(window, self.INDEX, 1)
      if output_watermark <= window.max_timestamp():
        nonspeculative_index = -1
        timing = windowed_value.PaneInfoTiming.EARLY
        if state.get_state(window, self.NONSPECULATIVE_INDEX):
          nonspeculative_index = state.get_state(
              window, self.NONSPECULATIVE_INDEX)
          state.add_state(window, self.NONSPECULATIVE_INDEX, 1)
          _LOGGER.warning(
              'Watermark moved backwards in time '
              'or late data moved window end forward.')
      else:
        nonspeculative_index = state.get_state(window, self.NONSPECULATIVE_INDEX)
        state.add_state(window, self.NONSPECULATIVE_INDEX, 1)
        timing = (
            windowed_value.PaneInfoTiming.ON_TIME if maybe_ontime and
            nonspeculative_index == 0 else windowed_value.PaneInfoTiming.LATE)
      pane_info = windowed_value.PaneInfo(
          index == 0, finished, timing, index, nonspeculative_index)
  
      values = state.get_state(window, self.ELEMENTS)
      if finished:
        # TODO(robertwb): allowed lateness
        state.clear_state(window, self.ELEMENTS)
        state.add_state(window, self.TOMBSTONE, 1)
      elif self.accumulation_mode == AccumulationMode.DISCARDING:
        state.clear_state(window, self.ELEMENTS)
  
      timestamp = state.get_state(window, self.WATERMARK_HOLD)
      if timestamp is None:
        # If no watermark hold was set, output at end of window.
        timestamp = window.max_timestamp()
      elif output_watermark < window.end and self.trigger_fn.has_ontime_pane():
        # Hold the watermark in case there is an empty pane that needs to be fired
        # at the end of the window.
        pass
      else:
        state.clear_state(window, self.WATERMARK_HOLD)
  
      return WindowedValue(values, timestamp, (window, ), pane_info)
  
  
  class InMemoryUnmergedState(UnmergedState):
    """In-memory implementation of UnmergedState.
  
    Used for batch and testing.
    """
    def __init__(self, defensive_copy=True):
      # TODO(robertwb): Skip defensive_copy in production if it's too expensive.
      self.timers = collections.defaultdict(dict)
      self.state = collections.defaultdict(lambda: collections.defaultdict(list))
      self.global_state = {}
      self.defensive_copy = defensive_copy
  
    def copy(self):
      cloned_object = InMemoryUnmergedState(defensive_copy=self.defensive_copy)
      cloned_object.timers = copy.deepcopy(self.timers)
      cloned_object.global_state = copy.deepcopy(self.global_state)
      for window in self.state:
        for tag in self.state[window]:
          cloned_object.state[window][tag] = copy.copy(self.state[window][tag])
      return cloned_object
  
    def set_global_state(self, tag, value):
      assert isinstance(tag, _ValueStateTag)
      if self.defensive_copy:
        value = copy.deepcopy(value)
      self.global_state[tag.tag] = value
  
    def get_global_state(self, tag, default=None):
      return self.global_state.get(tag.tag, default)
  
    def set_timer(self, window, name, time_domain, timestamp):
      self.timers[window][(name, time_domain)] = timestamp
  
    def clear_timer(self, window, name, time_domain):
      self.timers[window].pop((name, time_domain), None)
      if not self.timers[window]:
        del self.timers[window]
  
    def get_window(self, window_id):
      return window_id
  
    def add_state(self, window, tag, value):
      if self.defensive_copy:
        value = copy.deepcopy(value)
      if isinstance(tag, _ValueStateTag):
        self.state[window][tag.tag] = value
      elif isinstance(tag, _CombiningValueStateTag):
        # TODO(robertwb): Store merged accumulators.
        self.state[window][tag.tag].append(value)
      elif isinstance(tag, _ListStateTag):
        self.state[window][tag.tag].append(value)
      elif isinstance(tag, _SetStateTag):
        self.state[window][tag.tag].append(value)
      elif isinstance(tag, _WatermarkHoldStateTag):
        self.state[window][tag.tag].append(value)
      else:
        raise ValueError('Invalid tag.', tag)
  
    def get_state(self, window, tag):
      values = self.state[window][tag.tag]
      if isinstance(tag, _ValueStateTag):
        return values
      elif isinstance(tag, _CombiningValueStateTag):
        return tag.combine_fn.apply(values)
      elif isinstance(tag, _ListStateTag):
        return values
      elif isinstance(tag, _SetStateTag):
        return values
      elif isinstance(tag, _WatermarkHoldStateTag):
        return tag.timestamp_combiner_impl.combine_all(values)
      else:
        raise ValueError('Invalid tag.', tag)
  
    def clear_state(self, window, tag):
      self.state[window].pop(tag.tag, None)
      if not self.state[window]:
        self.state.pop(window, None)
  
    def get_timers(
        self, clear=False, watermark=MAX_TIMESTAMP, processing_time=None):
      """Gets expired timers and reports if there
      are any realtime timers set per state.
  
      Expiration is measured against the watermark for event-time timers,
      and against a wall clock for processing-time timers.
      """
      expired = []
      has_realtime_timer = False
      for window, timers in list(self.timers.items()):
        for (name, time_domain), timestamp in list(timers.items()):
          if time_domain == TimeDomain.REAL_TIME:
            time_marker = processing_time
            has_realtime_timer = True
          elif time_domain == TimeDomain.WATERMARK:
            time_marker = watermark
          else:
            _LOGGER.error(
                'TimeDomain error: No timers defined for time domain %s.',
                time_domain)
          if timestamp <= time_marker:
            expired.append((window, (name, time_domain, timestamp)))
            if clear:
              del timers[(name, time_domain)]
        if not timers and clear:
          del self.timers[window]
      return expired, has_realtime_timer
  
    def get_and_clear_timers(self, watermark=MAX_TIMESTAMP):
      return self.get_timers(clear=True, watermark=watermark)[0]
  
    def get_earliest_hold(self):
      earliest_hold = MAX_TIMESTAMP
      for unused_window, tagged_states in iteritems(self.state):
        # TODO(BEAM-2519): currently, this assumes that the watermark hold tag is
        # named "watermark".  This is currently only true because the only place
        # watermark holds are set is in the GeneralTriggerDriver, where we use
        # this name.  We should fix this by allowing enumeration of the tag types
        # used in adding state.
        if 'watermark' in tagged_states and tagged_states['watermark']:
          hold = min(tagged_states['watermark']) - TIME_GRANULARITY
          earliest_hold = min(earliest_hold, hold)
      return earliest_hold
  
    def __repr__(self):
      state_str = '\n'.join(
          '%s: %s' % (key, dict(state)) for key, state in self.state.items())
      return 'timers: %s\nstate: %s' % (dict(self.timers), state_str)
  