# HG changeset patch
# User Mike Frysinger <vapier@gentoo.org>
# Date 1354166686 18000
# Node ID 94901c23fe58ca227aacadcdc5d7ad3045171b24
# Parent  a9c9b2f9ec964e6750dd6f9a96767fe2e9201d85
rework dynamic tag walking to avoid doing it twice on the first call to iter_tags

Based on suggestions from David James

diff -r a9c9b2f9ec96 -r 94901c23fe58 elftools/elf/dynamic.py
--- a/elftools/elf/dynamic.py	Mon Nov 12 16:27:19 2012 -0500
+++ b/elftools/elf/dynamic.py	Thu Nov 29 00:24:46 2012 -0500
@@ -6,6 +6,8 @@
 # Mike Frysinger (vapier@gentoo.org)
 # This code is in the public domain
 #-------------------------------------------------------------------------------
+import itertools
+
 from .sections import Section
 from .segments import Segment
 from ..common.utils import struct_parse
@@ -20,11 +22,12 @@
         Similarly to Section objects, allows dictionary-like access to the
         dynamic tag.
     """
+
+    _HANDLED_TAGS = frozenset(['DT_NEEDED', 'DT_RPATH', 'DT_RUNPATH'])
+
     def __init__(self, entry, elffile):
         self.entry = entry
-        if entry.d_tag == 'DT_NEEDED' or \
-           entry.d_tag == 'DT_RPATH' or \
-           entry.d_tag == 'DT_RUNPATH':
+        if entry.d_tag in self._HANDLED_TAGS:
             dynstr = elffile.get_section_by_name('.dynstr')
             setattr(self, entry.d_tag[3:].lower(), dynstr.get_string(self.entry.d_val))
 
@@ -37,9 +40,7 @@
         return '<DynamicTag (%s): %r>' % (self.entry.d_tag, self.entry)
 
     def __str__(self):
-        if self.entry.d_tag == 'DT_NEEDED' or \
-           self.entry.d_tag == 'DT_RPATH' or \
-           self.entry.d_tag == 'DT_RUNPATH':
+        if self.entry.d_tag in self._HANDLED_TAGS:
             s = '"%s"' % getattr(self, self.entry.d_tag[3:].lower())
         else:
             s = '%#x' % self.entry.d_ptr
@@ -58,16 +59,12 @@
     def iter_tags(self, type=None):
         """ Yield all tags (limit to |type| if specified)
         """
-        offset = self._offset - self._tagsize
-        for i in range(self.num_tags):
-            offset += self._tagsize
-            entry = struct_parse(
-                self._elfstructs.Elf_Dyn,
-                self._stream,
-                stream_pos=offset)
-            if type is not None and entry.d_tag != type:
-                continue
-            yield DynamicTag(entry, self._elffile)
+        for n in itertools.count():
+            tag = self.get_tag(n)
+            if type is None or tag.entry.d_tag == type:
+                yield tag
+            if tag.entry.d_tag == 'DT_NULL':
+                break
 
     def get_tag(self, n):
         """ Get the tag at index #n from the file (DynamicTag object)
@@ -86,17 +83,11 @@
         if self._num_tags != -1:
             return self._num_tags
 
-        offset = self._offset
-        while True:
-            entry = struct_parse(
-                self._elfstructs.Elf_Dyn,
-                self._stream,
-                stream_pos=offset)
-            if entry.d_tag == 'DT_NULL':
-                self._num_tags = ((offset - self._offset) // self._tagsize) + 1
-                break
-            offset += self._tagsize
-        return self._num_tags
+        for n in itertools.count():
+            tag = self.get_tag(n)
+            if tag.entry.d_tag == 'DT_NULL':
+                self._num_tags = n
+                return n
 
 
 class DynamicSection(Section, Dynamic):