diff --git a/src/python/pants/java/nailgun_protocol.py b/src/python/pants/java/nailgun_protocol.py index 2220e82c145..a8198ab9375 100644 --- a/src/python/pants/java/nailgun_protocol.py +++ b/src/python/pants/java/nailgun_protocol.py @@ -85,9 +85,18 @@ class TruncatedPayloadError(TruncatedRead): """Raised if there is a socket error while reading the payload bytes.""" @classmethod - def _decode_unicode_seq(cls, seq): + def _encode_unicode(cls, payload, encoding='utf-8'): + """Convert unicode to bytes.""" + if isinstance(payload, six.text_type): + return payload.encode(encoding) + elif not isinstance(payload, six.binary_type): + raise TypeError('cannot encode type: {}'.format(type(payload))) + return payload + + @classmethod + def _encode_unicode_seq(cls, seq): for item in seq: - yield item.decode('utf-8') + yield cls._encode_unicode(item) @classmethod def send_request(cls, sock, working_dir, command, *arguments, **environment): @@ -98,7 +107,7 @@ def send_request(cls, sock, working_dir, command, *arguments, **environment): for item_tuple in environment.items(): cls.write_chunk(sock, ChunkType.ENVIRONMENT, - cls.ENVIRON_SEP.join(cls._decode_unicode_seq(item_tuple))) + cls.ENVIRON_SEP.join(cls._encode_unicode_seq(item_tuple))) cls.write_chunk(sock, ChunkType.WORKING_DIR, working_dir) cls.write_chunk(sock, ChunkType.COMMAND, command) @@ -144,11 +153,7 @@ def write_chunk(cls, sock, chunk_type, payload=b''): @classmethod def construct_chunk(cls, chunk_type, payload, encoding='utf-8'): """Construct and return a single chunk.""" - if isinstance(payload, six.text_type): - payload = payload.encode(encoding) - elif not isinstance(payload, six.binary_type): - raise TypeError('cannot encode type: {}'.format(type(payload))) - + payload = cls._encode_unicode(payload) header = struct.pack(cls.HEADER_FMT, len(payload), chunk_type) return header + payload