<s> from __future__ import absolute_import <EOL> import abc <EOL> import logging <EOL> import re <EOL> import six <EOL> from kafka . errors import IllegalStateError <EOL> from kafka . protocol . offset import OffsetResetStrategy <EOL> from kafka . structs import OffsetAndMetadata <EOL> log = logging . getLogger ( __name__ ) <EOL> class SubscriptionState ( object ) : <EOL> """<STR_LIT>""" <EOL> _SUBSCRIPTION_EXCEPTION_MESSAGE = ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> def __init__ ( self , offset_reset_strategy = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> offset_reset_strategy = getattr ( OffsetResetStrategy , <EOL> offset_reset_strategy . upper ( ) ) <EOL> except AttributeError : <EOL> log . warning ( '<STR_LIT>' ) <EOL> offset_reset_strategy = OffsetResetStrategy . NONE <EOL> self . _default_offset_reset_strategy = offset_reset_strategy <EOL> self . subscription = None <EOL> self . subscribed_pattern = None <EOL> self . _group_subscription = set ( ) <EOL> self . _user_assignment = set ( ) <EOL> self . assignment = dict ( ) <EOL> self . needs_partition_assignment = False <EOL> self . listener = None <EOL> self . needs_fetch_committed_offsets = True <EOL> def subscribe ( self , topics = ( ) , pattern = None , listener = None ) : <EOL> """<STR_LIT>""" <EOL> if self . _user_assignment or ( topics and pattern ) : <EOL> raise IllegalStateError ( self . _SUBSCRIPTION_EXCEPTION_MESSAGE ) <EOL> assert topics or pattern , '<STR_LIT>' <EOL> if pattern : <EOL> log . info ( '<STR_LIT>' , pattern ) <EOL> self . subscription = set ( ) <EOL> self . subscribed_pattern = re . compile ( pattern ) <EOL> else : <EOL> self . change_subscription ( topics ) <EOL> if listener and not isinstance ( listener , ConsumerRebalanceListener ) : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> self . listener = listener <EOL> def change_subscription ( self , topics ) : <EOL> """<STR_LIT>""" <EOL> if self . _user_assignment : <EOL> raise IllegalStateError ( self . _SUBSCRIPTION_EXCEPTION_MESSAGE ) <EOL> if self . subscription == set ( topics ) : <EOL> log . warning ( "<STR_LIT>" , <EOL> topics ) <EOL> return <EOL> log . info ( '<STR_LIT>' , topics ) <EOL> self . subscription = set ( topics ) <EOL> self . _group_subscription . update ( topics ) <EOL> self . needs_partition_assignment = True <EOL> for tp in set ( self . assignment . keys ( ) ) : <EOL> if tp . topic not in self . subscription : <EOL> del self . assignment [ tp ] <EOL> def group_subscribe ( self , topics ) : <EOL> """<STR_LIT>""" <EOL> if self . _user_assignment : <EOL> raise IllegalStateError ( self . _SUBSCRIPTION_EXCEPTION_MESSAGE ) <EOL> self . _group_subscription . update ( topics ) <EOL> def mark_for_reassignment ( self ) : <EOL> if self . _user_assignment : <EOL> raise IllegalStateError ( self . _SUBSCRIPTION_EXCEPTION_MESSAGE ) <EOL> assert self . subscription is not None , '<STR_LIT>' <EOL> self . _group_subscription . intersection_update ( self . subscription ) <EOL> self . needs_partition_assignment = True <EOL> def assign_from_user ( self , partitions ) : <EOL> """<STR_LIT>""" <EOL> if self . subscription is not None : <EOL> raise IllegalStateError ( self . _SUBSCRIPTION_EXCEPTION_MESSAGE ) <EOL> self . _user_assignment . clear ( ) <EOL> self . _user_assignment . update ( partitions ) <EOL> for partition in partitions : <EOL> if partition not in self . assignment : <EOL> self . _add_assigned_partition ( partition ) <EOL> for tp in set ( self . assignment . keys ( ) ) - self . _user_assignment : <EOL> del self . assignment [ tp ] <EOL> self . needs_partition_assignment = False <EOL> def assign_from_subscribed ( self , assignments ) : <EOL> """<STR_LIT>""" <EOL> if self . subscription is None : <EOL> raise IllegalStateError ( self . _SUBSCRIPTION_EXCEPTION_MESSAGE ) <EOL> for tp in assignments : <EOL> if tp . topic not in self . subscription : <EOL> raise ValueError ( "<STR_LIT>" % tp ) <EOL> self . assignment . clear ( ) <EOL> for tp in assignments : <EOL> self . _add_assigned_partition ( tp ) <EOL> self . needs_partition_assignment = False <EOL> log . info ( "<STR_LIT>" , assignments ) <EOL> def unsubscribe ( self ) : <EOL> """<STR_LIT>""" <EOL> self . subscription = None <EOL> self . _user_assignment . clear ( ) <EOL> self . assignment . clear ( ) <EOL> self . needs_partition_assignment = True <EOL> self . subscribed_pattern = None <EOL> def group_subscription ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _group_subscription <EOL> def seek ( self , partition , offset ) : <EOL> """<STR_LIT>""" <EOL> self . assignment [ partition ] . seek ( offset ) <EOL> def assigned_partitions ( self ) : <EOL> """<STR_LIT>""" <EOL> return set ( self . assignment . keys ( ) ) <EOL> def paused_partitions ( self ) : <EOL> """<STR_LIT>""" <EOL> return set ( partition for partition in self . assignment <EOL> if self . is_paused ( partition ) ) <EOL> def fetchable_partitions ( self ) : <EOL> """<STR_LIT>""" <EOL> fetchable = set ( ) <EOL> for partition , state in six . iteritems ( self . assignment ) : <EOL> if state . is_fetchable ( ) : <EOL> fetchable . add ( partition ) <EOL> return fetchable <EOL> def partitions_auto_assigned ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . subscription is not None <EOL> def all_consumed_offsets ( self ) : <EOL> """<STR_LIT>""" <EOL> all_consumed = { } <EOL> for partition , state in six . iteritems ( self . assignment ) : <EOL> if state . has_valid_position : <EOL> all_consumed [ partition ] = OffsetAndMetadata ( state . position , '<STR_LIT>' ) <EOL> return all_consumed <EOL> def need_offset_reset ( self , partition , offset_reset_strategy = None ) : <EOL> """<STR_LIT>""" <EOL> if offset_reset_strategy is None : <EOL> offset_reset_strategy = self . _default_offset_reset_strategy <EOL> self . assignment [ partition ] . await_reset ( offset_reset_strategy ) <EOL> def has_default_offset_reset_policy ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _default_offset_reset_strategy != OffsetResetStrategy . NONE <EOL> def is_offset_reset_needed ( self , partition ) : <EOL> return self . assignment [ partition ] . awaiting_reset <EOL> def has_all_fetch_positions ( self ) : <EOL> for state in self . assignment . values ( ) : <EOL> if not state . has_valid_position : <EOL> return False <EOL> return True <EOL> def missing_fetch_positions ( self ) : <EOL> missing = set ( ) <EOL> for partition , state in six . iteritems ( self . assignment ) : <EOL> if not state . has_valid_position : <EOL> missing . add ( partition ) <EOL> return missing <EOL> def is_assigned ( self , partition ) : <EOL> return partition in self . assignment <EOL> def is_paused ( self , partition ) : <EOL> return partition in self . assignment and self . assignment [ partition ] . paused <EOL> def is_fetchable ( self , partition ) : <EOL> return partition in self . assignment and self . assignment [ partition ] . is_fetchable ( ) <EOL> def pause ( self , partition ) : <EOL> self . assignment [ partition ] . pause ( ) <EOL> def resume ( self , partition ) : <EOL> self . assignment [ partition ] . resume ( ) <EOL> def _add_assigned_partition ( self , partition ) : <EOL> self . assignment [ partition ] = TopicPartitionState ( ) <EOL> class TopicPartitionState ( object ) : <EOL> def __init__ ( self ) : <EOL> self . committed = None <EOL> self . has_valid_position = False <EOL> self . paused = False <EOL> self . awaiting_reset = False <EOL> self . reset_strategy = None <EOL> self . _position = None <EOL> self . highwater = None <EOL> def _set_position ( self , offset ) : <EOL> assert self . has_valid_position , '<STR_LIT>' <EOL> self . _position = offset <EOL> def _get_position ( self ) : <EOL> return self . _position <EOL> position = property ( _get_position , _set_position , None , "<STR_LIT>" ) <EOL> def await_reset ( self , strategy ) : <EOL> self . awaiting_reset = True <EOL> self . reset_strategy = strategy <EOL> self . _position = None <EOL> self . has_valid_position = False <EOL> def seek ( self , offset ) : <EOL> self . _position = offset <EOL> self . awaiting_reset = False <EOL> self . reset_strategy = None <EOL> self . has_valid_position = True <EOL> def pause ( self ) : <EOL> self . paused = True <EOL> def resume ( self ) : <EOL> self . paused = False <EOL> def is_fetchable ( self ) : <EOL> return not self . paused and self . has_valid_position <EOL> class ConsumerRebalanceListener ( object ) : <EOL> """<STR_LIT>""" <EOL> __metaclass__ = abc . ABCMeta <EOL> @ abc . abstractmethod <EOL> def on_partitions_revoked ( self , revoked ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> @ abc . abstractmethod <EOL> def on_partitions_assigned ( self , assigned ) : <EOL> """<STR_LIT>""" <EOL> pass </s>
