diff --git a/lbrynet/dht/kbucket.py b/lbrynet/dht/kbucket.py index bb4cfc0dc..dfd3f5ae8 100644 --- a/lbrynet/dht/kbucket.py +++ b/lbrynet/dht/kbucket.py @@ -1,6 +1,10 @@ +import logging import constants +from distance import Distance from error import BucketFull +log = logging.getLogger(__name__) + class KBucket(object): """ Description - later @@ -28,7 +32,7 @@ class KBucket(object): already @param contact: The contact to add - @type contact: kademlia.contact.Contact + @type contact: dht.contact._Contact """ if contact in self._contacts: # Move the existing contact to the end of the list @@ -56,19 +60,21 @@ class KBucket(object): return contact raise IndexError(contactID) - def getContacts(self, count=-1, excludeContact=None): + def getContacts(self, count=-1, excludeContact=None, sort_distance_to=None): """ Returns a list containing up to the first count number of contacts @param count: The amount of contacts to return (if 0 or less, return all contacts) @type count: int - @param excludeContact: A contact to exclude; if this contact is in + @param excludeContact: A node id to exclude; if this contact is in the list of returned values, it will be discarded before returning. If a C{str} is passed as this argument, it must be the contact's ID. - @type excludeContact: kademlia.contact.Contact or str + @type excludeContact: str + @param sort_distance_to: Sort distance to the id, defaulting to the parent node id. If False don't + sort the contacts @raise IndexError: If the number of requested contacts is too large @@ -76,38 +82,35 @@ class KBucket(object): If no contacts are present an empty is returned @rtype: list """ + contacts = [contact for contact in self._contacts if contact.id != excludeContact] + # Return all contacts in bucket if count <= 0: - count = len(self._contacts) + count = len(contacts) # Get current contact number - currentLen = len(self._contacts) + currentLen = len(contacts) # If count greater than k - return only k contacts if count > constants.k: count = constants.k - # Check if count value in range and, - # if count number of contacts are available if not currentLen: - contactList = list() + return contacts - # length of list less than requested amount - elif currentLen < count: - contactList = self._contacts[0:currentLen] - # enough contacts in list + if sort_distance_to is False: + pass else: - contactList = self._contacts[0:count] + sort_distance_to = sort_distance_to or self._node_id + contacts.sort(key=lambda c: Distance(sort_distance_to)(c.id)) - if excludeContact in contactList: - contactList.remove(excludeContact) + return contacts[:min(currentLen, count)] def getBadOrUnknownContacts(self): contacts = self.getContacts(sort_distance_to=False) results = [contact for contact in contacts if contact.contact_is_good is False] results.extend(contact for contact in contacts if contact.contact_is_good is None) return results - return contactList def removeContact(self, contact): """ Remove the contact from the bucket