Merge pull request #49 from alexander255/master

Add `multiaddr.Multiaddr.split(self, maxsplit=-1)` and `multiaddr.Multiaddr.join(*addrs)`

Author: mhchia

multiaddr/multiaddr.py

@@ -6,6 +6,7 @@
     collections.abc = collections
 
 import six
+import varint
 
 from . import exceptions, protocols
 
@@ -33,7 +34,7 @@ def __getitem__(self, idx):
     __hash__ = collections.abc.KeysView._hash
 
     def __iter__(self):
-        for proto, _, _ in bytes_iter(self._mapping.to_bytes()):
+        for _, proto, _, _ in bytes_iter(self._mapping.to_bytes()):
             yield proto
 
 
@@ -52,7 +53,7 @@ def __getitem__(self, idx):
         raise IndexError("Protocol item list index out of range")
 
     def __iter__(self):
-        for proto, codec, part in bytes_iter(self._mapping.to_bytes()):
+        for _, proto, codec, part in bytes_iter(self._mapping.to_bytes()):
             if codec.SIZE != 0:
                 try:
                     # If we have an address, return it
@@ -130,6 +131,12 @@ def __init__(self, addr):
         else:
             raise TypeError("MultiAddr must be bytes, str or another MultiAddr instance")
 
+    @classmethod
+    def join(cls, *addrs):
+        """Concatenate the values of the given MultiAddr strings or objects,
+        encapsulating each successive MultiAddr value with the previous ones."""
+        return cls(b"".join(map(lambda a: cls(a).to_bytes(), addrs)))
+
     def __eq__(self, other):
         """Checks if two Multiaddr objects are exactly equal."""
         return self._bytes == other._bytes
@@ -171,6 +178,29 @@ def protocols(self):
         """Returns a list of Protocols this Multiaddr includes."""
         return MultiAddrKeys(self)
 
+    def split(self, maxsplit=-1):
+        """Returns the list of individual path components this MultiAddr is made
+        up of."""
+        final_split_offset = -1
+        results = []
+        for idx, (offset, proto, codec, part_value) in enumerate(bytes_iter(self._bytes)):
+            # Split at most `maxplit` times
+            if idx == maxsplit:
+                final_split_offset = offset
+                break
+
+            # Re-assemble binary MultiAddr representation
+            part_size = varint.encode(len(part_value)) if codec.SIZE < 0 else b""
+            part = b"".join((proto.vcode, part_size, part_value))
+
+            # Add MultiAddr with the given value
+            results.append(self.__class__(part))
+        # Add final item with remainder of MultiAddr if there is anything left
+        if final_split_offset >= 0:
+            results.append(self.__class__(self._bytes[final_split_offset:]))
+
+        return results
+
     keys = protocols
 
     def items(self):
@@ -185,9 +215,7 @@ def encapsulate(self, other):
         For example:
             /ip4/1.2.3.4 encapsulate /tcp/80 = /ip4/1.2.3.4/tcp/80
         """
-        mb = self.to_bytes()
-        ob = Multiaddr(other).to_bytes()
-        return Multiaddr(b''.join([mb, ob]))
+        return self.join(self, other)
 
     def decapsulate(self, other):
         """Remove a Multiaddr wrapping.
@@ -209,7 +237,7 @@ def value_for_protocol(self, proto):
 
         Returns
         -------
-        union[object, NoneType]
+        Union[object, NoneType]
             The parsed protocol value for the given protocol code or ``None``
             if the given protocol does not require any value
 

multiaddr/protocols.py

@@ -1,5 +1,4 @@
 # -*- coding: utf-8 -*-
-import binascii
 import six
 import varint
 
@@ -118,31 +117,6 @@ def __repr__(self):
         )
 
 
-def _uvarint(buf):
-    """Reads a varint from a bytes buffer and returns the value and # bytes"""
-    x = 0
-    s = 0
-    for i, b_str in enumerate(buf):
-        if six.PY3:
-            b = b_str
-        else:  # pragma: no cover (PY2)
-            b = int(binascii.b2a_hex(b_str), 16)
-        if b < 0x80:
-            if i > 9 or (i == 9 and b > 1):
-                # Code 34 apparently means `OverflowError`
-                # (as opposed to other `ArithmeticError` types)
-                raise OverflowError(34, "UVarInt too large")
-            return (x | b << s, i + 1)
-        x |= (b & 0x7f) << s
-        s += 7
-    return 0, 0
-
-
-def read_varint_code(buf):
-    num, n = _uvarint(buf)
-    return int(num), n
-
-
 # Protocols is the list of multiaddr protocols supported by this module.
 PROTOCOLS = [
     Protocol(P_IP4, 'ip4', 'ip4'),

multiaddr/transforms.py

@@ -1,4 +1,5 @@
 # -*- encoding: utf-8 -*-
+import io
 import six
 import varint
 
@@ -9,7 +10,6 @@
 
 from .protocols import protocol_with_code
 from .protocols import protocol_with_name
-from .protocols import read_varint_code
 
 
 def string_to_bytes(string):
@@ -29,7 +29,7 @@ def string_to_bytes(string):
 
 def bytes_to_string(buf):
     st = [u'']  # start with empty string so we get a leading slash on join()
-    for proto, codec, part in bytes_iter(buf):
+    for _, proto, codec, part in bytes_iter(buf):
         st.append(proto.name)
         if codec.SIZE != 0:
             try:
@@ -43,11 +43,11 @@ def bytes_to_string(buf):
     return u'/'.join(st)
 
 
-def size_for_addr(codec, buf):
+def size_for_addr(codec, buf_io):
     if codec.SIZE >= 0:
-        return codec.SIZE // 8, 0
+        return codec.SIZE // 8
     else:
-        return read_varint_code(buf)
+        return varint.decode_stream(buf_io)
 
 
 def string_iter(string):
@@ -82,8 +82,10 @@ def string_iter(string):
 
 
 def bytes_iter(buf):
-    while buf:
-        code, num_bytes_read = read_varint_code(buf)
+    buf_io = io.BytesIO(buf)
+    while buf_io.tell() < len(buf):
+        offset = buf_io.tell()
+        code = varint.decode_stream(buf_io)
         proto = None
         try:
             proto = protocol_with_code(code)
@@ -97,7 +99,5 @@ def bytes_iter(buf):
                 ),
                 exc,
             )
-        size, num_bytes_read2 = size_for_addr(codec, buf[num_bytes_read:])
-        length = size + num_bytes_read2 + num_bytes_read
-        yield proto, codec, buf[(length - size):length]
-        buf = buf[length:]
+        size = size_for_addr(codec, buf_io)
+        yield offset, proto, codec, buf_io.read(size)

tests/test_multiaddr.py

@@ -159,6 +159,33 @@ def test_invalid_protocols_with_string(proto_string):
         protocols_with_string(proto_string)
 
 
+@pytest.mark.parametrize(
+    'proto_string,maxsplit,expected',
+    [("/ip4/1.2.3.4", -1, ("/ip4/1.2.3.4",)),
+     ("/ip4/0.0.0.0", 0, ("/ip4/0.0.0.0",)),
+     ("/ip6/::1", 1, ("/ip6/::1",)),
+     ("/onion/timaq4ygg2iegci7:80/http", 0, ("/onion/timaq4ygg2iegci7:80/http",)),
+     ("/ip4/127.0.0.1/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC/tcp/1234", 1,
+      ("/ip4/127.0.0.1", "/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC/tcp/1234",)),
+     ("/ip4/1.2.3.4/tcp/80/unix/a/b/c/d/e/f", -1,
+      ("/ip4/1.2.3.4", "/tcp/80", "/unix/a/b/c/d/e/f"))])
+def test_split(proto_string, maxsplit, expected):
+    assert tuple(map(str, Multiaddr(proto_string).split(maxsplit))) == expected
+
+
+@pytest.mark.parametrize(
+    'proto_parts,expected',
+    [(("/ip4/1.2.3.4",), "/ip4/1.2.3.4"),
+     ((b"\x04\x00\x00\x00\x00",), "/ip4/0.0.0.0"),
+     (("/ip6/::1",), "/ip6/::1"),
+     (("/onion/timaq4ygg2iegci7:80/http",), "/onion/timaq4ygg2iegci7:80/http"),
+     ((b"\x04\x7F\x00\x00\x01", "/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC/tcp/1234",),
+      "/ip4/127.0.0.1/p2p/QmcgpsyWgH8Y8ajJz1Cu72KnS5uo2Aa2LpzU7kinSupNKC/tcp/1234"),
+     (("/ip4/1.2.3.4", "/tcp/80", "/unix/a/b/c/d/e/f"), "/ip4/1.2.3.4/tcp/80/unix/a/b/c/d/e/f")])
+def test_join(proto_parts, expected):
+    assert str(Multiaddr.join(*proto_parts)) == expected
+
+
 def test_encapsulate():
     m1 = Multiaddr("/ip4/127.0.0.1/udp/1234")
     m2 = Multiaddr("/udp/5678")

tests/test_protocols.py

@@ -161,15 +161,3 @@ def test_add_protocol_twice(patch_protocols, valid_params):
 def test_protocol_repr():
     proto = protocols.protocol_with_name('ip4')
     assert "Protocol(code=4, name='ip4', codec='ip4')" == repr(proto)
-
-
-@pytest.mark.parametrize("buf", [
-    b'\x90\x91\x92\x93\x94\x95\x96\x97\x98\x99\x01',
-    b'\x90\x91\x92\x93\x94\x95\x96\x97\x98\x02'])
-def test_overflowing_varint(buf):
-    with pytest.raises(OverflowError):
-        protocols.read_varint_code(buf)
-
-
-def test_nonterminated_varint():
-    assert protocols.read_varint_code(b'\x80\x80') == (0, 0)

tests/test_transforms.py

@@ -1,4 +1,6 @@
 # -*- encoding: utf-8 -*-
+import io
+
 import pytest
 
 from multiaddr.codecs import codec_by_name
@@ -57,7 +59,8 @@
     ('p2p', b'\x40\x50\x60\x51', (64, 1)),
 ])
 def test_size_for_addr(codec_name, buf, expected):
-    assert size_for_addr(codec_by_name(codec_name), buf) == expected
+    buf_io = io.BytesIO(buf)
+    assert (size_for_addr(codec_by_name(codec_name), buf_io), buf_io.tell()) == expected
 
 
 @pytest.mark.parametrize("buf, expected", [
@@ -69,7 +72,7 @@ def test_size_for_addr(codec_name, buf, expected):
       (_names_to_protocols["tcp"], b'\x10\xe1')]),
 ])
 def test_bytes_iter(buf, expected):
-    assert list((proto, val) for proto, _, val in bytes_iter(buf)) == expected
+    assert list((proto, val) for _, proto, _, val in bytes_iter(buf)) == expected
 
 
 @pytest.mark.parametrize("proto, buf, expected", ADDR_BYTES_MAP_STR_TEST_DATA)