Unify multiaddr parsing to two functions: codec.{bytes,string}_iter

Author: ntninja

multiaddr/codec.py

@@ -32,55 +32,26 @@ def find_codec_by_name(name):
 def string_to_bytes(string):
     if not string:
         return b''
-    # consume trailing slashes
-    if not string.startswith('/'):
-        raise ValueError("invalid multiaddr, must begin with /")
-    string = string.rstrip('/')
-    sp = string.split('/')
-
-    # skip the first element, since it starts with /
-    sp.pop(0)
+
     bs = []
-    while sp:
-        element = sp.pop(0)
-        proto = protocol_with_name(element)
+    for proto, codec, value in string_iter(string):
         bs.append(varint.encode(proto.code))
-        try:
-            codec = find_codec_by_name(proto.codec)
-        except ImportError as exc:
-            six.raise_from(ValueError("failed to parse %s addr: unknown" % proto.name), exc)
-        if codec.SIZE == 0:
-            continue
-        if len(sp) < 1:
-            raise ValueError(
-                "protocol requires address, none given: %s" % proto.name)
-        if codec.IS_PATH:
-            sp = ["/" + "/".join(sp)]
-        bs.append(codec.to_bytes(proto, sp.pop(0)))
+        if value is not None:
+            bs.append(codec.to_bytes(proto, value))
     return b''.join(bs)
 
 
 def bytes_to_string(buf):
-    st = ['']  # start with empty string so we get a leading slash on join()
-    while buf:
-        maddr_component = ""
-        code, num_bytes_read = read_varint_code(buf)
-        buf = buf[num_bytes_read:]
-        proto = protocol_with_code(code)
-        maddr_component += proto.name
-        try:
-            codec = find_codec_by_name(proto.codec)
-        except ImportError as exc:
-            six.raise_from(ValueError("failed to parse %s addr: unknown" % proto.name), exc)
-        size = size_for_addr(codec, buf)
-        if size > 0:
-            addr = codec.to_string(proto, buf[:size])
-            if not (codec.IS_PATH and addr[0] == '/'):
-                maddr_component += '/'
-            maddr_component += addr
-        st.append(maddr_component)
-        buf = buf[size:]
-    return '/'.join(st)
+	st = [u'']  # start with empty string so we get a leading slash on join()
+	for proto, codec, part in bytes_iter(buf):
+		st.append(proto.name)
+		if codec.SIZE != 0:
+			value = codec.to_string(proto, part)
+			if codec.IS_PATH and value[0] == u'/':
+				st.append(value[1:])
+			else:
+				st.append(value)
+	return u'/'.join(st)
 
 
 def size_for_addr(codec, buf):
@@ -91,17 +62,47 @@ def size_for_addr(codec, buf):
         return size + num_bytes_read
 
 
-def bytes_split(buf):
-    ret = []
-    while buf:
-        code, num_bytes_read = read_varint_code(buf)
-        proto = protocol_with_code(code)
-        try:
-            codec = find_codec_by_name(proto.codec)
-        except ImportError as exc:
-            six.raise_from(ValueError("failed to parse %s addr: unknown" % proto.name), exc)
-        size = size_for_addr(codec, buf[num_bytes_read:])
-        length = size + num_bytes_read
-        ret.append(buf[:length])
-        buf = buf[length:]
-    return ret
+def string_iter(string):
+	if not string.startswith(u'/'):
+		raise ValueError("invalid multiaddr, must begin with /")
+	# consume trailing slashes
+	string = string.rstrip(u'/')
+	sp = string.split(u'/')
+
+	# skip the first element, since it starts with /
+	sp.pop(0)
+	while sp:
+		element = sp.pop(0)
+		proto = protocol_with_name(element)
+		try:
+			codec = find_codec_by_name(proto.codec)
+		except ImportError as exc:
+			six.raise_from(ValueError("failed to parse %s addr: unknown" % proto.name), exc)
+		value = None
+		if codec.SIZE != 0:
+			if len(sp) < 1:
+				raise ValueError(
+					"protocol requires address, none given: %s" % proto.name)
+			if codec.IS_PATH:
+				value = "/" + "/".join(sp)
+				if not six.PY2:
+					sp.clear()
+				else:
+					sp = []
+			else:
+				value = sp.pop(0)
+		yield proto, codec, value
+
+
+def bytes_iter(buf):
+	while buf:
+		code, num_bytes_read = read_varint_code(buf)
+		proto = protocol_with_code(code)
+		try:
+			codec = find_codec_by_name(proto.codec)
+		except ImportError as exc:
+			six.raise_from(ValueError("failed to parse %s addr: unknown" % proto.name), exc)
+		size = size_for_addr(codec, buf[num_bytes_read:])
+		length = size + num_bytes_read
+		yield proto, codec, buf[num_bytes_read:length]
+		buf = buf[length:]

multiaddr/multiaddr.py

@@ -3,12 +3,9 @@
 
 import six
 
-from .codec import find_codec_by_name
-from .codec import size_for_addr
+from .codec import bytes_iter
 from .codec import string_to_bytes
 from .codec import bytes_to_string
-from .protocols import protocol_with_code
-from .protocols import read_varint_code
 
 
 class ProtocolNotFoundException(Exception):
@@ -89,20 +86,7 @@ def to_bytes(self):
 
     def protocols(self):
         """Returns a list of Protocols this Multiaddr includes."""
-        buf = self.to_bytes()
-        protos = []
-        while buf:
-            code, num_bytes_read = read_varint_code(buf)
-            proto = protocol_with_code(code)
-            try:
-                codec = find_codec_by_name(proto.codec)
-            except ImportError as exc:
-                six.raise_from(ValueError("failed to parse %s addr: unknown" % proto.name), exc)
-            protos.append(proto)
-            buf = buf[num_bytes_read:]
-            size = size_for_addr(codec, buf)
-            buf = buf[size:]
-        return protos
+        return list(proto for proto, _, _ in bytes_iter(self.to_bytes()))
 
     def encapsulate(self, other):
         """Wrap this Multiaddr around another.
@@ -136,24 +120,15 @@ def decapsulate(self, other):
 
     def value_for_protocol(self, code):
         """Return the value (if any) following the specified protocol."""
-        from .util import split
-
         if not isinstance(code, int):
             raise ValueError("code type should be `int`, code={}".format(code))
 
-        for sub_addr in split(self):
-            protocol = sub_addr.protocols()[0]
-            if protocol.code == code:
-                # e.g. if `sub_addr=/unix/123`, then `addr_parts=['', 'unix', '123']`
-                addr_parts = str(sub_addr).split("/")
-                if protocol.path:
-                    return "/" + "/".join(addr_parts[2:])
-                if len(addr_parts) > 3:
-                    raise ValueError("Unknown Protocol format")
-                elif len(addr_parts) == 3:
+        for proto, codec, part in bytes_iter(self.to_bytes()):
+            if proto.code == code:
+                if codec.SIZE != 0:
                     # If we have an address, return it
-                    return addr_parts[2]
-                elif len(addr_parts) == 2:
+                    return codec.to_string(proto, part)
+                else:
                     # We were given something like '/utp', which doesn't have
                     # an address, so return ''
                     return ''

multiaddr/util.py

@@ -1,26 +1,9 @@
 import six
 import struct
 
-from .codec import bytes_split
 from .multiaddr import Multiaddr
 
 
-def split(ma):
-    """Return the sub-address portions of a multiaddr"""
-    addrs = []
-    bb = bytes_split(ma.to_bytes())
-    for addr in bb:
-        addrs.append(Multiaddr(addr))
-    return addrs
-
-
-def join(multiaddrs):
-    bs = []
-    for ma in multiaddrs:
-        bs.append(ma.to_bytes())
-    return Multiaddr(b''.join(bs))
-
-
 if hasattr(int, 'from_bytes'):
     def packed_net_bytes_to_int(b):
         """Convert the given big-endian byte-string to an int."""

tests/test_codec.py

@@ -2,7 +2,7 @@
 import pytest
 
 from multiaddr.codec import find_codec_by_name
-from multiaddr.codec import bytes_split
+from multiaddr.codec import bytes_iter
 from multiaddr.codec import bytes_to_string
 from multiaddr.codec import size_for_addr
 from multiaddr.codec import string_to_bytes
@@ -59,13 +59,13 @@ def test_size_for_addr(codec_name, buf, expected):
 @pytest.mark.parametrize("buf, expected", [
     # "/ip4/127.0.0.1/udp/1234/ip4/127.0.0.1/tcp/4321"
     (b'\x04\x7f\x00\x00\x01\x91\x02\x04\xd2\x04\x7f\x00\x00\x01\x06\x10\xe1',
-     [b'\x04\x7f\x00\x00\x01',
-      b'\x91\x02\04\xd2',
-      b'\x04\x7f\x00\x00\x01',
-      b'\x06\x10\xe1']),
+     [(_names_to_protocols["ip4"], b'\x7f\x00\x00\x01'),
+      (_names_to_protocols["udp"], b'\x04\xd2'),
+      (_names_to_protocols["ip4"], b'\x7f\x00\x00\x01'),
+      (_names_to_protocols["tcp"], b'\x10\xe1')]),
 ])
-def test_bytes_split(buf, expected):
-    assert bytes_split(buf) == expected
+def test_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)

tests/test_multiaddr.py

@@ -16,8 +16,6 @@
 from multiaddr.protocols import P_TCP
 from multiaddr.protocols import P_UDP
 from multiaddr.protocols import P_UNIX
-from multiaddr.util import split
-from multiaddr.util import join
 
 
 @pytest.mark.parametrize(
@@ -117,23 +115,6 @@ def test_eq():
     assert m4 == m3
 
 
-@pytest.mark.parametrize(
-    'test_vals',
-    [("/ip4/1.2.3.4/udp/1234", ["/ip4/1.2.3.4", "/udp/1234"]),
-     ("/ip4/1.2.3.4/tcp/1/ip4/2.3.4.5/udp/2",
-      ["/ip4/1.2.3.4", "/tcp/1", "/ip4/2.3.4.5", "/udp/2"]),
-     ("/ip4/1.2.3.4/utp/ip4/2.3.4.5/udp/2/udt",
-      ["/ip4/1.2.3.4", "/utp", "/ip4/2.3.4.5", "/udp/2", "/udt"])])
-def test_bytes_split_and_join(test_vals):
-    string, expected = test_vals
-    mm = Multiaddr(string)
-    split_m = split(mm)
-    for i, addr in enumerate(split_m):
-        assert str(addr) == expected[i]
-    joined = join(split_m)
-    assert mm == joined
-
-
 def test_protocols():
     ma = Multiaddr("/ip4/127.0.0.1/udp/1234")
     protos = ma.protocols()
@@ -246,21 +227,6 @@ def test_bad_initialization_wrong_type():
         Multiaddr(42)
 
 
-def test_get_value_too_many_fields_protocol(monkeypatch):
-    """
-    This test patches the Multiaddr's string representation to return
-    an invalid string in order to test that value_for_protocol properly
-    throws a ValueError.  This avoids some of the error checking in
-    the constructor and is easier to patch, thus the actual values
-    that the constructor specifies is ignored by the test.
-    """
-    monkeypatch.setattr("multiaddr.multiaddr.Multiaddr.__str__",
-                        lambda ignore: str('/udp/1234/5678'))
-    a = Multiaddr("/ip4/127.0.0.1/udp/1234")
-    with pytest.raises(ValueError):
-        a.value_for_protocol(P_UDP)
-
-
 def test_value_for_protocol_argument_wrong_type():
     a = Multiaddr("/ip4/127.0.0.1/udp/1234")
     with pytest.raises(ValueError):