[tor-commits] [stem/master] Extend _read_until_keyword to read until multiple keywords

[atagar at torproject.org]

commit 1b26a46ed3c2b26293474d9349efb247f5888a65
Author: Ravi Chandra Padmala <neenaoffline at gmail.com>
Date:   Wed Aug 8 08:33:48 2012 +0530

    Extend _read_until_keyword to read until multiple keywords
---
 stem/descriptor/__init__.py             |   11 ++++++-----
 stem/descriptor/extrainfo_descriptor.py |    4 ++--
 stem/descriptor/server_descriptor.py    |    6 +++---
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/stem/descriptor/__init__.py b/stem/descriptor/__init__.py
index 8193380..b1f3ab6 100644
--- a/stem/descriptor/__init__.py
+++ b/stem/descriptor/__init__.py
@@ -148,19 +148,20 @@ class Descriptor(object):
   def __str__(self):
     return self._raw_contents
 
-def _read_until_keyword(keyword, descriptor_file, inclusive = False):
+def _read_until_keywords(keywords, descriptor_file, inclusive = False):
   """
-  Reads from the descriptor file until we get to the given keyword or reach the
+  Reads from the descriptor file until we get to one of the given keywords or reach the
   end of the file.
   
-  :param str keyword: keyword we want to read until
+  :param str,list keywords: keyword(s) we want to read until
   :param file descriptor_file: file with the descriptor content
   :param bool inclusive: includes the line with the keyword if True
   
-  :returns: list with the lines until we find the keyword
+  :returns: list with the lines until we find one of the keywords
   """
   
   content = []
+  if type(keywords) == str: keywords = (keywords,)
   
   while True:
     last_position = descriptor_file.tell()
@@ -170,7 +171,7 @@ def _read_until_keyword(keyword, descriptor_file, inclusive = False):
     if " " in line: line_keyword = line.split(" ", 1)[0]
     else: line_keyword = line.strip()
     
-    if line_keyword == keyword:
+    if line_keyword in keywords:
       if inclusive: content.append(line)
       else: descriptor_file.seek(last_position)
       
diff --git a/stem/descriptor/extrainfo_descriptor.py b/stem/descriptor/extrainfo_descriptor.py
index a589139..8d181ae 100644
--- a/stem/descriptor/extrainfo_descriptor.py
+++ b/stem/descriptor/extrainfo_descriptor.py
@@ -131,11 +131,11 @@ def parse_file(descriptor_file, validate = True):
   """
   
   while True:
-    extrainfo_content = stem.descriptor._read_until_keyword("router-signature", descriptor_file)
+    extrainfo_content = stem.descriptor._read_until_keywords("router-signature", descriptor_file)
     
     # we've reached the 'router-signature', now include the pgp style block
     block_end_prefix = stem.descriptor.PGP_BLOCK_END.split(' ', 1)[0]
-    extrainfo_content += stem.descriptor._read_until_keyword(block_end_prefix, descriptor_file, True)
+    extrainfo_content += stem.descriptor._read_until_keywords(block_end_prefix, descriptor_file, True)
     
     if extrainfo_content:
       yield RelayExtraInfoDescriptor("".join(extrainfo_content), validate)
diff --git a/stem/descriptor/server_descriptor.py b/stem/descriptor/server_descriptor.py
index 3dc6216..9b8ef0d 100644
--- a/stem/descriptor/server_descriptor.py
+++ b/stem/descriptor/server_descriptor.py
@@ -109,12 +109,12 @@ def parse_file(descriptor_file, validate = True):
   # to the caller).
   
   while True:
-    annotations = stem.descriptor._read_until_keyword("router", descriptor_file)
-    descriptor_content = stem.descriptor._read_until_keyword("router-signature", descriptor_file)
+    annotations = stem.descriptor._read_until_keywords("router", descriptor_file)
+    descriptor_content = stem.descriptor._read_until_keywords("router-signature", descriptor_file)
     
     # we've reached the 'router-signature', now include the pgp style block
     block_end_prefix = stem.descriptor.PGP_BLOCK_END.split(' ', 1)[0]
-    descriptor_content += stem.descriptor._read_until_keyword(block_end_prefix, descriptor_file, True)
+    descriptor_content += stem.descriptor._read_until_keywords(block_end_prefix, descriptor_file, True)
     
     if descriptor_content:
       # strip newlines from annotations