Merge pull request #30 from amathael/master

Created wrap_many(*args) function

Author: Hrabal

tempy/tempy.py

@@ -2,19 +2,20 @@
 # @author: Federico Cerchiari <[email protected]>
 """Main Tempy classes"""
 import html
+from collections import deque, Iterable
 from copy import copy
-from uuid import uuid4
-from itertools import chain
 from functools import wraps
-from collections import deque
+from itertools import chain
 from types import GeneratorType
+from uuid import uuid4
 
-from .exceptions import TagError, WrongContentError, DOMModByKeyError, DOMModByIndexError
+from .exceptions import TagError, WrongContentError, WrongArgsError, DOMModByKeyError, DOMModByIndexError
 from .tempyrepr import REPRFinder
 
 
 class DOMGroup:
     """Wrapper used to manage element insertion."""
+
     def __init__(self, name, obj):
         super().__init__()
         if not name and issubclass(obj.__class__, DOMElement):
@@ -167,7 +168,7 @@ def __imul__(self, n):
         if n == 0:
             self.parent.pop(self._own_index)
             return self
-        return self.after(self * (n-1))
+        return self.after(self * (n - 1))
 
     def to_code(self, pretty=False):
         ret = []
@@ -252,6 +253,7 @@ def _iter_child_renders(self, pretty=False):
                     # this trick is used to avoid circular imports
                     class Patched(tempyREPR_cls, DOMElement):
                         pass
+
                     child = Patched(child)
             try:
                 yield child.render(pretty=pretty)
@@ -272,14 +274,17 @@ def content_receiver(reverse=False):
         Takes args and kwargs and calls the decorated method one time for each argument provided.
         The reverse parameter should be used for prepending (relative to self) methods.
         """
+
         def _receiver(func):
             @wraps(func)
             def wrapped(inst, *tags, **kwtags):
                 for i, dom_group in yield_domgroups(tags, kwtags, reverse):
                     inst._stable = False
                     func(inst, i, dom_group)
                 return inst
+
             return wrapped
+
         return _receiver
 
     def _insert(self, dom_group, idx=None, prepend=False):
@@ -298,7 +303,7 @@ def _insert(self, dom_group, idx=None, prepend=False):
             for i_group, elem in enumerate(dom_group):
                 if elem is not None:
                     # Element insertion in this DOMElement childs
-                    self.childs.insert(idx+i_group, elem)
+                    self.childs.insert(idx + i_group, elem)
                     # Managing child attributes if needed
                     if issubclass(elem.__class__, DOMElement):
                         elem.parent = self
@@ -393,7 +398,6 @@ def append_to(self, father):
 
     def wrap(self, other):
         """Wraps this element inside another empty tag."""
-        # TODO: make multiple with content_receiver
         if other.childs:
             raise TagError(self, 'Wrapping in a non empty Tag is forbidden.')
         if self.parent:
@@ -402,6 +406,51 @@ def wrap(self, other):
         other.append(self)
         return self
 
+    def wrap_many(self, *args, strict=False):
+        """Wraps different copies of this element inside all empty tags
+        listed in params or param's (non-empty) iterators.
+
+        Returns list of copies of this element wrapped inside args
+        or None if not succeeded, in the same order and same structure,
+        i.e. args = (Div(), (Div())) -> value = (A(...), (A(...)))
+
+        If on some args it must raise TagError, it will only if strict is True,
+        otherwise it will do nothing with them and return Nones on their positions"""
+
+        for arg in args:
+            is_elem = arg and isinstance(arg, DOMElement)
+            is_elem_iter = (not is_elem and arg and isinstance(arg, Iterable) and
+                            isinstance(iter(arg).__next__(), DOMElement))
+            if not (is_elem or is_elem_iter):
+                raise WrongArgsError(self, 'Argument {} is not DOMElement nor iterable of DOMElements'.format(arg))
+
+        wcopies = []
+        failure = []
+
+        def wrap_next(tag, idx):
+            nonlocal wcopies, failure
+            next_copy = self.__copy__()
+            try:
+                return next_copy.wrap(tag)
+            except TagError:
+                failure.append(idx)
+                return next_copy
+
+        for arg_idx, arg in enumerate(args):
+            if isinstance(arg, DOMElement):
+                wcopies.append(wrap_next(arg, (arg_idx, -1)))
+            else:
+                iter_wcopies = []
+                for iter_idx, t in enumerate(arg):
+                    iter_wcopies.append(wrap_next(t, (arg_idx, iter_idx)))
+                wcopies.append(type(arg)(iter_wcopies))
+
+        if failure and strict:
+            raise TagError(self, 'Wrapping in a non empty Tag is forbidden, failed on arguments ' +
+                           ', '.join(list(map(lambda idx: str(idx[0]) if idx[1] == -1 else '[{1}] of {0}'.format(*idx),
+                                              failure))))
+        return wcopies
+
     def wrap_inner(self, other):
         self.move_childs(other)
         self(other)
@@ -599,7 +648,6 @@ def render(self, *args, **kwargs):
 
 
 class Escaped(DOMElement):
-
     def __init__(self, content, **kwargs):
         super().__init__(**kwargs)
         self._render = content

tests/test_DOMElement.py

@@ -4,14 +4,13 @@
 """
 import unittest
 
+from tempy.elements import Tag, TagAttrs
+from tempy.exceptions import WrongContentError, WrongArgsError, TagError, DOMModByKeyError, DOMModByIndexError
 from tempy.tags import Div, A, P, Html, Head, Body
 from tempy.tempy import DOMElement, DOMGroup, Escaped
-from tempy.elements import Tag, TagAttrs
-from tempy.exceptions import WrongContentError, TagError, DOMModByKeyError, DOMModByIndexError
 
 
 class TestDOMelement(unittest.TestCase):
-
     def setUp(self):
         self.page = Html()
 
@@ -84,13 +83,13 @@ def test_after(self):
         new1 = Div().append_to(self.page)
         new2 = Div()
         new1.after(new2)
-        self.assertEqual(new1._own_index, new2._own_index-1)
+        self.assertEqual(new1._own_index, new2._own_index - 1)
 
     def test_before(self):
         new1 = Div().append_to(self.page)
         new2 = Div()
         new1.before(new2)
-        self.assertEqual(new1._own_index, new2._own_index+1)
+        self.assertEqual(new1._own_index, new2._own_index + 1)
 
     def test_prepend(self):
         self.page(Div(), Div())
@@ -134,6 +133,53 @@ def test_wrap(self):
         self.assertTrue(to_wrap in container)
         self.assertTrue(container in outermost)
 
+    def test_wrap_many(self):
+        def flatten(cnt):
+            res = []
+            for el in cnt:
+                if isinstance(el, DOMElement):
+                    res.append(el)
+                else:
+                    res.extend(el)
+            return res
+
+        def test_return_values(inp, outp):
+            self.assertEqual(len(inp), len(outp))
+            for _ in range(len(inp)):
+                t1, t2 = type(inp[_]), type(outp[_])
+                self.assertTrue(t1 == t2 or
+                                issubclass(t1, DOMElement) and issubclass(t2, DOMElement))
+
+        def test_correctly_wrapped(child, parent):
+            self.assertTrue(child in parent)
+            self.assertTrue(child.get_parent() == parent)
+
+        # check if it works correct with correct arguments
+        args = (Div(), [Div(), Div()], (Div(), Div()))
+        new = A().wrap_many(*args)
+        test_return_values(args, new)
+        for c, p in zip(flatten(new), flatten(args)):
+            test_correctly_wrapped(c, p)
+
+        # check if it raises TagError with strict and returns None without
+        args = (Div()(A()), (Div(), Div()))
+        with self.assertRaisesRegex(TagError, r'^.+arguments 0$'):
+            A().wrap_many(*args, strict=True)
+        new = A().wrap_many(*args)
+        self.assertIs(new[0].get_parent(), None)
+
+        args = (Div()(A()), (Div(), Div()(A())))
+        with self.assertRaisesRegex(TagError, r'^.+arguments 0, \[1\] of 1'):
+            A().wrap_many(*args, strict=True)
+        new = A().wrap_many(*args)
+        self.assertIs(new[0].get_parent(), None)
+        self.assertIs(new[1][1].get_parent(), None)
+
+        # check if it raises WrongArgsError
+        args = (Div(), '')
+        with self.assertRaises(WrongArgsError):
+            A().wrap_many(*args)
+
     def test_replace_with(self):
         old = Div().append_to(self.page)
         old.replace_with(A())