netstring: add required_trailer= argument

This commit is contained in:
Brian Warner 2008-09-26 09:57:54 -07:00
parent 27a30b90d7
commit 98c8e25709
2 changed files with 16 additions and 5 deletions

View File

@ -12,6 +12,10 @@ class Netstring(unittest.TestCase):
self.failUnlessRaises(ValueError, split_netstring, a, 3) self.failUnlessRaises(ValueError, split_netstring, a, 3)
self.failUnlessRaises(ValueError, split_netstring, a+" extra", 2) self.failUnlessRaises(ValueError, split_netstring, a+" extra", 2)
self.failUnlessRaises(ValueError, split_netstring, a+" extra", 2, False) self.failUnlessRaises(ValueError, split_netstring, a+" extra", 2, False)
self.failUnlessEqual(split_netstring(a+"++", 2, required_trailer="++"),
("hello", "world"))
self.failUnlessRaises(ValueError,
split_netstring, a+"+", 2, required_trailer="not")
def test_extra(self): def test_extra(self):
a = netstring("hello") a = netstring("hello")

View File

@ -4,11 +4,18 @@ def netstring(s):
assert isinstance(s, str), s # no unicode here assert isinstance(s, str), s # no unicode here
return "%d:%s," % (len(s), s,) return "%d:%s," % (len(s), s,)
def split_netstring(data, numstrings, allow_leftover=False): def split_netstring(data, numstrings,
allow_leftover=False,
required_trailer=""):
"""like string.split(), but extracts netstrings. If allow_leftover=False, """like string.split(), but extracts netstrings. If allow_leftover=False,
returns numstrings elements, and throws ValueError if there was leftover I return numstrings elements, and throw ValueError if there was leftover
data. If allow_leftover=True, returns numstrings+1 elements, in which the data that does not exactly equal 'required_trailer'. If
last element is the leftover data (possibly an empty string)""" allow_leftover=True, required_trailer must be empty, and I return
numstrings+1 elements, in which the last element is the leftover data
(possibly an empty string)"""
assert not (allow_leftover and required_trailer)
elements = [] elements = []
assert numstrings >= 0 assert numstrings >= 0
while data: while data:
@ -25,7 +32,7 @@ def split_netstring(data, numstrings, allow_leftover=False):
raise ValueError("ran out of netstrings") raise ValueError("ran out of netstrings")
if allow_leftover: if allow_leftover:
return tuple(elements + [data]) return tuple(elements + [data])
if data: if data != required_trailer:
raise ValueError("leftover data in netstrings") raise ValueError("leftover data in netstrings")
return tuple(elements) return tuple(elements)