Commits

Gael Pasgrimaud committed cf3908a

fix wrap / add wrapAll

Comments (0)

Files changed (2)

pyquery/pyquery.py

         for tag in self:
             wrapper = deepcopy(value)
             # FIXME: using iterchildren is probably not optimal
-            childs = [c for c in wrapper.iterchildren()]
-            if not childs:
-                child = wrapper
+            if not wrapper.getchildren():
+                wrapper.append(deepcopy(tag))
             else:
+                childs = [c for c in wrapper.iterchildren()]
                 child = childs[-1]
-            child.append(deepcopy(tag))
-            nodes.append(child)
+                child.append(deepcopy(tag))
+            nodes.append(wrapper)
 
             parent = tag.getparent()
             if parent is not None:
-                index = parent.index(tag)
-                # FIXME: using iterchildren is probably not optimal
                 for t in parent.iterchildren():
                     if t is tag:
-                        t.addnext(child)
+                        t.addnext(wrapper)
                         parent.remove(t)
                         break
         self[:] = nodes
         return self
 
+    def wrapAll(self, value):
+        """Wrap all the elements in the matched set into a single wrapper element::
+
+            >>> d = PyQuery('<div><span>Hey</span><span>you !</span></div>')
+            >>> print d('span').wrapAll('<div id="wrapper"></div>')
+            <div><div id="wrapper"><span>Hey</span><span>you !</span></div></div>
+
+        """
+        if not self:
+            return self
+
+        assert isinstance(value, basestring)
+        value = fromstring(value)
+        wrapper = deepcopy(value)
+        if not wrapper.getchildren():
+            child = wrapper
+        else:
+            childs = [c for c in wrapper.iterchildren()]
+            child = childs[-1]
+
+        replace_childs = True
+        parent = self[0].getparent()
+        if parent is None:
+            parent = no_default
+
+        # add nodes to wrapper and check parent
+        for tag in self:
+            child.append(deepcopy(tag))
+            if tag.getparent() is not parent:
+                replace_childs = False
+
+        # replace nodes i parent if possible
+        if parent is not no_default and replace_childs:
+            childs = [c for c in parent.iterchildren()]
+            if len(childs) == len(self):
+                for tag in self:
+                    parent.remove(tag)
+                parent.append(wrapper)
+
+        self[:] = [wrapper]
+        return self
+
     def replaceWith(self, value):
         """replace nodes by value
         """
         >>> s = d('span')
         >>> s is d
         False
-        >>> s.wrap('<div id="wrapper"></div>')
-        [<div#wrapper>]
+        >>> s.wrap('<div><div id="wrapper"></div></div>')
+        [<div>]
 
     We get the original doc with new node::
 
         >>> print d
-        <div id="bouh"><div id="wrapper"><span>youhou</span></div></div>
+        <div id="bouh"><div><div id="wrapper"><span>youhou</span></div></div></div>
+
+    Complex wrapAll::
+
+        >>> doc = pq('<div><span>Hey</span><span>you !</span></div>')
+        >>> s = doc('span')
+        >>> s.wrapAll('<div id="wrapper"></div>')
+        [<div#wrapper>]
+
+        >>> print doc
+        <div><div id="wrapper"><span>Hey</span><span>you !</span></div></div>
     """
 
 class DocTest(doctest.DocFileCase):
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.