[Lxml-checkins] r42062 - in lxml/trunk: . src/lxml src/lxml/tests

[scoder at codespeak.net]

Author: scoder
Date: Sun Apr 15 11:43:31 2007
New Revision: 42062

Modified:
   lxml/trunk/CHANGES.txt
   lxml/trunk/src/lxml/etree.pyx
   lxml/trunk/src/lxml/tests/test_etree.py
Log:
support for element.attrib.pop()

Modified: lxml/trunk/CHANGES.txt
==============================================================================
--- lxml/trunk/CHANGES.txt	(original)
+++ lxml/trunk/CHANGES.txt	Sun Apr 15 11:43:31 2007
@@ -8,6 +8,8 @@
 Features added
 --------------
 
+* Element.attrib now has a ``pop()`` method
+
 * Extended type annotation in objectify: cleaner annotation namespace setup
   plus new ``xsiannotate()`` and ``deannotate()`` functions
 

Modified: lxml/trunk/src/lxml/etree.pyx
==============================================================================
--- lxml/trunk/src/lxml/etree.pyx	(original)
+++ lxml/trunk/src/lxml/etree.pyx	Sun Apr 15 11:43:31 2007
@@ -1411,6 +1411,20 @@
         for key, value in sequence_or_dict:
             _setAttributeValue(self._element, key, value)
 
+    def pop(self, key, *default):
+        if python.PyTuple_GET_SIZE(default) > 1:
+            raise TypeError, "pop expected at most 2 arguments, got %d" % \
+                  (python.PyTuple_GET_SIZE(default)+1)
+        result = _getAttributeValue(self._element, key, None)
+        if result is None:
+            if python.PyTuple_GET_SIZE(default) == 0:
+                raise KeyError, key
+            else:
+                return python.PyTuple_GET_ITEM(default, 0)
+        else:
+            _delAttribute(self._element, key)
+            return result
+
     # ACCESSORS
     def __repr__(self):
         return repr(dict( _attributeIteratorFactory(self._element, 3) ))

Modified: lxml/trunk/src/lxml/tests/test_etree.py
==============================================================================
--- lxml/trunk/src/lxml/tests/test_etree.py	(original)
+++ lxml/trunk/src/lxml/tests/test_etree.py	Sun Apr 15 11:43:31 2007
@@ -70,6 +70,39 @@
         self.assertEquals("TEST", root.get("attr"))
         self.assertRaises(TypeError, root.set, "newattr", 5)
 
+    def test_attrib_pop(self):
+        ElementTree = self.etree.ElementTree
+        
+        f = StringIO('<doc one="One" two="Two"/>')
+        doc = ElementTree(file=f)
+        root = doc.getroot()
+        self.assertEquals('One', root.attrib['one'])
+        self.assertEquals('Two', root.attrib['two'])
+
+        self.assertEquals('One', root.attrib.pop('one'))
+
+        self.assertEquals(None, root.attrib.get('one'))
+        self.assertEquals('Two', root.attrib['two'])
+
+    def test_attrib_pop_unknown(self):
+        root = self.etree.XML('<doc one="One" two="Two"/>')
+        self.assertRaises(KeyError, root.attrib.pop, 'NONE')
+
+        self.assertEquals('One', root.attrib['one'])
+        self.assertEquals('Two', root.attrib['two'])
+
+    def test_attrib_pop_default(self):
+        root = self.etree.XML('<doc one="One" two="Two"/>')
+        self.assertEquals('Three', root.attrib.pop('three', 'Three'))
+
+    def test_attrib_pop_empty_default(self):
+        root = self.etree.XML('<doc/>')
+        self.assertEquals('Three', root.attrib.pop('three', 'Three'))
+
+    def test_attrib_pop_invalid_args(self):
+        root = self.etree.XML('<doc one="One" two="Two"/>')
+        self.assertRaises(TypeError, root.attrib.pop, 'One', None, None)
+
     def test_pi(self):
         # lxml.etree separates target and text
         Element = self.etree.Element